From b1313a466a883f440976ac60c13805078f7d0ac6 Mon Sep 17 00:00:00 2001 From: ray Date: Mon, 29 Jun 2026 19:47:27 +0800 Subject: [PATCH 01/45] refactor: add quantizer and distance --- src/turbo/quantizer/distance.h | 107 +++++++++++++++++++ src/turbo/quantizer/quantizer.h | 180 ++++++++++++++++++++++++++++++++ 2 files changed, 287 insertions(+) create mode 100644 src/turbo/quantizer/distance.h create mode 100644 src/turbo/quantizer/quantizer.h diff --git a/src/turbo/quantizer/distance.h b/src/turbo/quantizer/distance.h new file mode 100644 index 000000000..bc8af6c1a --- /dev/null +++ b/src/turbo/quantizer/distance.h @@ -0,0 +1,107 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +namespace zvec { +namespace turbo { + +//! A callable distance handle bound to a quantized query vector. +//! +//! DistanceImpl owns the quantized query bytes and a dispatched +//! DistanceFunc. Invoking `operator()(candidate)` computes the distance +//! between the stored query and the given candidate vector, which is +//! expected to already be in the same quantized layout. +class DistanceImpl { + public: + DistanceImpl() = default; + + DistanceImpl(DistanceFunc func, std::string quantized_query, size_t dim) + : func_(std::move(func)), + query_storage_(std::move(quantized_query)), + dim_(dim) {} + + DistanceImpl(DistanceFunc func, BatchDistanceFunc batch_func, + std::string quantized_query, size_t dim) + : func_(std::move(func)), + batch_func_(std::move(batch_func)), + query_storage_(std::move(quantized_query)), + dim_(dim) {} + + //! Whether the handle is ready to compute distances. + bool valid() const { + return static_cast(func_); + } + + //! Whether a batch distance function is available. + bool batch_valid() const { + return static_cast(batch_func_); + } + + //! Compute the distance between the stored query and `candidate`. + float operator()(const void *candidate) const { + float d = 0.0f; + func_(candidate, query_storage_.data(), dim_, &d); + return d; + } + + //! Compute distances for a batch of `num` candidates against the + //! stored query. Falls back to the scalar path when no batch function + //! is bound. + void batch(const void **candidates, size_t num, float *out) const { + if (batch_func_) { + batch_func_(candidates, query_storage_.data(), num, dim_, out); + return; + } + for (size_t i = 0; i < num; ++i) { + out[i] = 0.0f; + func_(candidates[i], query_storage_.data(), dim_, out + i); + } + } + + //! Access the quantized query bytes (for pairwise helpers). + const std::string &query_storage() const { + return query_storage_; + } + + size_t dim() const { + return dim_; + } + + //! Raw scalar distance function (operates on already-quantized + //! candidates). Useful for pairwise node-vs-node distance where no + //! stored query is involved. + const DistanceFunc &func() const { + return func_; + } + + //! Raw batch distance function. + const BatchDistanceFunc &batch_func() const { + return batch_func_; + } + + private: + DistanceFunc func_{}; + BatchDistanceFunc batch_func_{}; + std::string query_storage_{}; + size_t dim_{0}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h new file mode 100644 index 000000000..c6d1db2bc --- /dev/null +++ b/src/turbo/quantizer/quantizer.h @@ -0,0 +1,180 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include "distance.h" + +using namespace zvec::core; + +namespace zvec { +namespace turbo { + +//! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. +constexpr uint32_t kQuantizerMagic = 0x52545A51u; +//! Current quantizer serialization format version. +constexpr uint16_t kQuantizerSerVersion = 1; + +//! Self-describing, fixed-size header that prefixes every serialized quantizer. +//! The type-specific payload (scalar params, codebook, rotation matrix, ...) +//! follows immediately after this header. +struct QuantizerSerHeader { + uint32_t magic; // kQuantizerMagic + uint16_t version; // kQuantizerSerVersion + uint16_t quant_type; // QuantizeType + uint32_t dim; // original dim (sanity check) + uint32_t metric; // MetricType (sanity check) + uint32_t payload_size; // bytes following the header + uint32_t reserved; // 0, for future use / alignment +}; +static_assert(sizeof(QuantizerSerHeader) == 24, + "QuantizerSerHeader must be 24 bytes"); + +class Quantizer { + public: + typedef std::shared_ptr Pointer; + + Quantizer() {} + virtual ~Quantizer() {} + + //! Initialize quantizer with index metadata and parameters + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) = 0; + + //! Get the output metadata after initialization + virtual const IndexMeta &meta() const = 0; + + //! Input data type accepted by the quantizer + virtual DataType input_data_type() const = 0; + + //! Data type + virtual QuantizeType type() const { + return type_; + } + + //! Dimensionality of the input vectors + virtual int dim() const = 0; + + //! Train the quantizer with a contiguous batch of data + virtual int train(const void * /*data*/, size_t /*num*/, size_t /*stride*/) { + return IndexError_NotImplemented; + } + + //! Whether the quantizer requires training before use + virtual bool require_train() const = 0; + + //! Train the quantizer with data from an IndexHolder + virtual int train(IndexHolder::Pointer /*holder*/) { + return IndexError_NotImplemented; + } + + //! Byte length of a quantized datapoint vector + virtual size_t quantized_datapoint_vector_length() const = 0; + + //! Byte length of a quantized query vector + virtual size_t quantized_query_vector_length() const = 0; + + //! Quantize a datapoint vector + virtual void quantize_data(const void *input, void *output) const = 0; + + //! Quantize a query vector + virtual void quantize_query(const void *input, void *output) const = 0; + + //! Distance between a quantized datapoint and a quantized query + virtual float calc_distance_dp_query(const void *dp, + const void *query) const = 0; + + //! Batched distance between quantized datapoints and a quantized query + virtual void calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const = 0; + + //! Distance between a quantized datapoint and an unquantized query + virtual float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const = 0; + + //! Batched distance between quantized datapoints and an unquantized query + virtual void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const = 0; + + //! Distance between two quantized datapoints + virtual float calc_distance_dp_dp(const void *dp1, const void *dp2) const = 0; + + //! Quantize a query vector for search + virtual int quantize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const { + return IndexError_NotImplemented; + } + + //! Dequantize a result vector back to original format + virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + virtual DistanceImpl distance(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/) const { + return DistanceImpl{}; + } + + //! Serialize quantizer parameters + virtual int serialize(std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + //! Deserialize quantizer parameters + virtual int deserialize(std::string & /*in*/) { + return IndexError_NotImplemented; + } + + //! Deserialize quantizer parameters from a raw, possibly mmap-backed buffer + //! (zero-copy entry point for large payloads such as codebooks/matrices). + virtual int deserialize(const void * /*data*/, size_t /*len*/) { + return IndexError_NotImplemented; + } + + protected: + //! Map a metric name (e.g. "SquaredEuclidean", "Cosine", + //! "InnerProduct", "MipsSquaredEuclidean") to its MetricType. + static MetricType metric_from_name(const std::string &name) { + if (name == "SquaredEuclidean") { + return MetricType::kSquaredEuclidean; + } + if (name == "Cosine") { + return MetricType::kCosine; + } + if (name == "InnerProduct") { + return MetricType::kInnerProduct; + } + if (name == "MipsSquaredEuclidean") { + return MetricType::kMipsSquaredEuclidean; + } + return MetricType::kUnknown; + } + + QuantizeType type_{QuantizeType::kDefault}; + uint32_t extra_meta_size_{0}; +}; + +} // namespace turbo +} // namespace zvec From 45d9b54a10259cfe988dfbadc39bc516dbef5d19 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 14:16:10 +0800 Subject: [PATCH 02/45] fix: add turo quantizer --- src/turbo/CMakeLists.txt | 4 +- .../record_quantized_int8/common.h | 0 .../record_quantized_int8/cosine.cc | 0 .../record_quantized_int8/cosine.h | 0 .../squared_euclidean.cc | 0 .../record_quantized_int8/squared_euclidean.h | 0 .../avx512_vnni/uniform_int8/quantize.cc | 0 .../avx512_vnni/uniform_int8/quantize.h | 0 .../uniform_int8/squared_euclidean.cc | 0 .../uniform_int8/squared_euclidean.h | 0 tests/turbo/turbo_fp32_quantizer_test.cc | 83 +++++++++++++++++++ 11 files changed, 85 insertions(+), 2 deletions(-) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/common.h (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/cosine.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/cosine.h (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/squared_euclidean.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/record_quantized_int8/squared_euclidean.h (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/quantize.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/quantize.h (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/squared_euclidean.cc (100%) rename src/turbo/{ => distance}/avx512_vnni/uniform_int8/squared_euclidean.h (100%) create mode 100644 tests/turbo/turbo_fp32_quantizer_test.cc diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index 9cbb2fac7..fcf9dc431 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -19,7 +19,7 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h) # subdirectory). if(NOT ANDROID AND AUTO_DETECT_ARCH) if (HOST_ARCH MATCHES "^(x86|x64)$") - file(GLOB_RECURSE AVX512_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/avx512_vnni/*.cc) + file(GLOB_RECURSE AVX512_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc) set_source_files_properties( ${AVX512_VNNI_SRCS} PROPERTIES @@ -32,5 +32,5 @@ cc_library( NAME zvec_turbo STATIC STRICT PACKED SRCS ${ALL_SRCS} LIBS zvec_ailego - INCS ${CMAKE_CURRENT_SOURCE_DIR} ${PROJECT_ROOT_DIR}/src/include + INCS ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/distance ${PROJECT_ROOT_DIR}/src/include ) diff --git a/src/turbo/avx512_vnni/record_quantized_int8/common.h b/src/turbo/distance/avx512_vnni/record_quantized_int8/common.h similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/common.h rename to src/turbo/distance/avx512_vnni/record_quantized_int8/common.h diff --git a/src/turbo/avx512_vnni/record_quantized_int8/cosine.cc b/src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.cc similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/cosine.cc rename to src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.cc diff --git a/src/turbo/avx512_vnni/record_quantized_int8/cosine.h b/src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.h similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/cosine.h rename to src/turbo/distance/avx512_vnni/record_quantized_int8/cosine.h diff --git a/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.cc similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc rename to src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.cc diff --git a/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.h b/src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.h similarity index 100% rename from src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.h rename to src/turbo/distance/avx512_vnni/record_quantized_int8/squared_euclidean.h diff --git a/src/turbo/avx512_vnni/uniform_int8/quantize.cc b/src/turbo/distance/avx512_vnni/uniform_int8/quantize.cc similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/quantize.cc rename to src/turbo/distance/avx512_vnni/uniform_int8/quantize.cc diff --git a/src/turbo/avx512_vnni/uniform_int8/quantize.h b/src/turbo/distance/avx512_vnni/uniform_int8/quantize.h similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/quantize.h rename to src/turbo/distance/avx512_vnni/uniform_int8/quantize.h diff --git a/src/turbo/avx512_vnni/uniform_int8/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.cc similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/squared_euclidean.cc rename to src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.cc diff --git a/src/turbo/avx512_vnni/uniform_int8/squared_euclidean.h b/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.h similarity index 100% rename from src/turbo/avx512_vnni/uniform_int8/squared_euclidean.h rename to src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.h diff --git a/tests/turbo/turbo_fp32_quantizer_test.cc b/tests/turbo/turbo_fp32_quantizer_test.cc new file mode 100644 index 000000000..40165a5d3 --- /dev/null +++ b/tests/turbo/turbo_fp32_quantizer_test.cc @@ -0,0 +1,83 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include "zvec/core/framework/index_factory.h" + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +TEST(Fp32Quantizer, General) { + std::mt19937 gen(15583); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t COUNT = 10000; + const size_t DIMENSION = 12; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Fp32Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0u, quantizer->init(meta, params)); + + auto holder = + std::make_shared>( + DIMENSION); + for (size_t i = 0; i < COUNT; ++i) { + zvec::ailego::NumericalVector vec(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + vec[j] = dist(gen); + } + holder->emplace(i + 1, vec); + } + EXPECT_EQ(COUNT, holder->count()); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, holder->data_type()); + + ASSERT_EQ(0u, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::string quant_buffer; + std::string dequant_buffer; + + for (; iter->is_valid(); iter->next()) { + EXPECT_TRUE(iter->data()); + + IndexQueryMeta qmeta; + quant_buffer.clear(); + EXPECT_EQ(0, quantizer->quantize( + iter->data(), + IndexQueryMeta(holder->data_type(), holder->dimension()), + &quant_buffer, &qmeta)); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, qmeta.data_type()); + EXPECT_EQ(holder->dimension(), qmeta.dimension()); + + dequant_buffer.clear(); + EXPECT_EQ( + 0, quantizer->dequantize(quant_buffer.data(), qmeta, &dequant_buffer)); + + const float *original_data = reinterpret_cast(iter->data()); + const float *dequantize_data = + reinterpret_cast(dequant_buffer.data()); + for (size_t i = 0; i < holder->dimension(); ++i) { + EXPECT_NEAR(original_data[i], dequantize_data[i], 1e-3); + } + } +} \ No newline at end of file From e82d73f343e75680ba7c66710d91f022c462dd3d Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 15:26:47 +0800 Subject: [PATCH 03/45] refactor: add quantizer --- src/core/framework/index_factory.cc | 13 ++++++++++++ .../zvec/core/framework/index_factory.h | 19 ++++++++++++++++++ src/include/zvec/turbo/turbo.h | 20 +++++++++++++++++++ tests/CMakeLists.txt | 1 + tests/turbo/CMakeLists.txt | 14 +++++++++++++ 5 files changed, 67 insertions(+) create mode 100644 tests/turbo/CMakeLists.txt diff --git a/src/core/framework/index_factory.cc b/src/core/framework/index_factory.cc index 69fe0e98d..e93f57bc7 100644 --- a/src/core/framework/index_factory.cc +++ b/src/core/framework/index_factory.cc @@ -257,5 +257,18 @@ std::vector IndexFactory::AllRefiners(void) { return ailego::Factory::Classes(); } +std::shared_ptr IndexFactory::CreateQuantizer( + const std::string &name) { + return ailego::Factory::MakeShared(name.c_str()); +} + +bool IndexFactory::HasQuantizer(const std::string &name) { + return ailego::Factory::Has(name.c_str()); +} + +std::vector IndexFactory::AllQuantizers(void) { + return ailego::Factory::Classes(); +} + } // namespace core } // namespace zvec diff --git a/src/include/zvec/core/framework/index_factory.h b/src/include/zvec/core/framework/index_factory.h index d891eaa5a..00e77894c 100644 --- a/src/include/zvec/core/framework/index_factory.h +++ b/src/include/zvec/core/framework/index_factory.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include #include @@ -167,6 +168,16 @@ struct IndexFactory { //! Retrieve all refiner classes static std::vector AllRefiners(void); + + //! Create a quantizer by name + static std::shared_ptr CreateQuantizer( + const std::string &name); + + //! Test if the quantizer exists + static bool HasQuantizer(const std::string &name); + + //! Retrieve all quantizer classes + static std::vector AllQuantizers(void); }; //! Register Index Metric @@ -283,5 +294,13 @@ struct IndexFactory { #define INDEX_FACTORY_REGISTER_REFINER(__IMPL__, ...) \ INDEX_FACTORY_REGISTER_REFINER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) +//! Register Quantizer +#define INDEX_FACTORY_REGISTER_QUANTIZER_ALIAS(__NAME__, __IMPL__, ...) \ + AILEGO_FACTORY_REGISTER(__NAME__, turbo::Quantizer, __IMPL__, ##__VA_ARGS__) + +//! Register Quantizer +#define INDEX_FACTORY_REGISTER_QUANTIZER(__IMPL__, ...) \ + INDEX_FACTORY_REGISTER_QUANTIZER_ALIAS(__IMPL__, __IMPL__, ##__VA_ARGS__) + } // namespace core } // namespace zvec diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 2fbf6d680..df353bded 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -36,18 +36,38 @@ using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale, enum class MetricType { kSquaredEuclidean, kCosine, + kInnerProduct, kMipsSquaredEuclidean, kUnknown, }; enum class DataType { + kInt4, kInt8, + kFp16, + kFp32, kUnknown, }; enum class QuantizeType { kDefault, kUniform, + kRecord, + kFp16, + kFp32, + kPQ, + kRabit +}; + +enum class CpuArchType { + kAuto, + kScalar, + kSSE, + kAVX, + kAVX2, + kAVX512, + kAVX512VNNI, + kAVX512FP16 }; DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7be2294dd..e3b54ee24 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -7,6 +7,7 @@ include_directories(${PROJECT_ROOT_DIR}) cc_directories(ailego) cc_directories(db) cc_directories(core) +cc_directories(turbo) if(BUILD_C_BINDINGS) cc_directories(c) endif() diff --git a/tests/turbo/CMakeLists.txt b/tests/turbo/CMakeLists.txt new file mode 100644 index 000000000..8a3527d41 --- /dev/null +++ b/tests/turbo/CMakeLists.txt @@ -0,0 +1,14 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) + +file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc) + +foreach(CC_SRCS ${ALL_TEST_SRCS}) + get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) + cc_gtest( + NAME ${CC_TARGET} + STRICT + LIBS zvec_ailego core_framework core_metric core_quantizer zvec_turbo + SRCS ${CC_SRCS} + INCS . ${PROJECT_ROOT_DIR}/src/core/ ${PROJECT_ROOT_DIR}/src/turbo/ + ) +endforeach() \ No newline at end of file From 829f53e198aa092f0a62ee6f6da2e56b11e4b894 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 15:29:46 +0800 Subject: [PATCH 04/45] refactor: update meta --- src/core/framework/index_meta.cc | 4 +- src/include/zvec/core/framework/index_meta.h | 50 +++++++++++++++++++- 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/src/core/framework/index_meta.cc b/src/core/framework/index_meta.cc index 11d54cb63..d0eadb02d 100644 --- a/src/core/framework/index_meta.cc +++ b/src/core/framework/index_meta.cc @@ -30,7 +30,8 @@ struct IndexMetaFormatHeader { uint32_t space_id; uint32_t attachment_offset; uint32_t attachment_size; - uint8_t reserved_[4092]; + uint32_t extra_meta_size; + uint8_t reserved_[4088]; }; static_assert(sizeof(IndexMetaFormatHeader) % 32 == 0, @@ -47,6 +48,7 @@ void IndexMeta::serialize(std::string *out) const { format.dimension = dimension_; format.unit_size = unit_size_; format.space_id = space_id_; + format.extra_meta_size = extra_meta_size_; if (!metric_name_.empty()) { ailego::Params item; diff --git a/src/include/zvec/core/framework/index_meta.h b/src/include/zvec/core/framework/index_meta.h index 3a09aaefb..c6c5137db 100644 --- a/src/include/zvec/core/framework/index_meta.h +++ b/src/include/zvec/core/framework/index_meta.h @@ -40,6 +40,7 @@ class IndexMeta { DT_BINARY64 = 8, }; + /*! Major Orders */ enum MajorOrder { @@ -77,6 +78,7 @@ class IndexMeta { dimension_(rhs.dimension_), unit_size_(rhs.unit_size_), element_size_(rhs.element_size_), + extra_meta_size_(rhs.extra_meta_size_), space_id_(rhs.space_id_), metric_revision_(rhs.metric_revision_), converter_revision_(rhs.converter_revision_), @@ -112,6 +114,7 @@ class IndexMeta { dimension_(rhs.dimension_), unit_size_(rhs.unit_size_), element_size_(rhs.element_size_), + extra_meta_size_(rhs.extra_meta_size_), space_id_(rhs.space_id_), metric_revision_(rhs.metric_revision_), converter_revision_(rhs.converter_revision_), @@ -173,6 +176,7 @@ class IndexMeta { searcher_params_ = std::move(rhs.searcher_params_); streamer_params_ = std::move(rhs.streamer_params_); attributes_ = std::move(rhs.attributes_); + extra_meta_size_ = rhs.extra_meta_size_; return *this; } @@ -211,6 +215,7 @@ class IndexMeta { searcher_params_ = std::move(rhs.searcher_params_); streamer_params_ = std::move(rhs.streamer_params_); attributes_ = std::move(rhs.attributes_); + extra_meta_size_ = rhs.extra_meta_size_; return *this; } @@ -249,6 +254,7 @@ class IndexMeta { searcher_params_.clear(); streamer_params_.clear(); attributes_.clear(); + extra_meta_size_ = 0; } //! Retrieve major order information @@ -281,6 +287,11 @@ class IndexMeta { return element_size_; } + //! Retrieve extra meta size in bytes + uint32_t extra_meta_size(void) const { + return extra_meta_size_; + } + //! Retrieve space id uint64_t space_id(void) const { return space_id_; @@ -451,6 +462,11 @@ class IndexMeta { this->set_meta(data_type, UnitSizeof(data_type), dim); } + //! Set extra meta size + void set_extra_meta_size(uint32_t size) { + extra_meta_size_ = size; + } + //! Set information of metric template void set_metric(TName &&name, uint32_t rev, TParams &¶ms) { @@ -586,6 +602,7 @@ class IndexMeta { uint32_t dimension_{0}; uint32_t unit_size_{0}; uint32_t element_size_{0}; + uint32_t extra_meta_size_{0}; uint64_t space_id_{0}; uint32_t metric_revision_{0}; uint32_t converter_revision_{0}; @@ -632,6 +649,19 @@ class IndexQueryMeta { unit_size_(unit), element_size_(IndexMeta::ElementSizeof(data_type, unit, dim)) {} + //! Constructor + IndexQueryMeta(IndexMeta::MetaType meta_type, IndexMeta::DataType data_type, + uint32_t unit, uint32_t dim, uint32_t quantize_type, + uint32_t extra_meta_size) + : meta_type_(meta_type), + data_type_(data_type), + dimension_(dim), + unit_size_(unit), + quantize_type_(quantize_type), + extra_meta_size_(extra_meta_size), + element_size_(IndexMeta::ElementSizeof(data_type, unit, dim) + + extra_meta_size_) {} + //! Constructor IndexQueryMeta(IndexMeta::DataType data_type, uint32_t dim) : IndexQueryMeta{IndexMeta::MetaType::MT_DENSE, data_type, @@ -676,7 +706,8 @@ class IndexQueryMeta { //! Set dimension of feature void set_dimension(uint32_t dim) { dimension_ = dim; - element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim); + element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dim) + + extra_meta_size_; } //! Set meta type @@ -694,7 +725,8 @@ class IndexQueryMeta { data_type_ = data_type; dimension_ = dim; unit_size_ = unit; - element_size_ = IndexMeta::ElementSizeof(data_type, unit, dim); + element_size_ = + IndexMeta::ElementSizeof(data_type, unit, dim) + extra_meta_size_; } //! Set meta information of feature @@ -702,11 +734,25 @@ class IndexQueryMeta { this->set_meta(data_type, IndexMeta::UnitSizeof(data_type), dim); } + //! Set meta information of feature with quantize type and extra meta size + void set_meta(IndexMeta::DataType data_type, uint32_t dim, + uint32_t quantize_type, uint32_t extra_meta_size) { + data_type_ = data_type; + dimension_ = dim; + unit_size_ = IndexMeta::UnitSizeof(data_type); + quantize_type_ = quantize_type; + extra_meta_size_ = extra_meta_size; + element_size_ = + IndexMeta::ElementSizeof(data_type, unit_size_, dim) + extra_meta_size_; + } + private: IndexMeta::MetaType meta_type_{IndexMeta::MetaType::MT_DENSE}; IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED}; uint32_t dimension_{0}; uint32_t unit_size_{0}; + uint32_t quantize_type_{0}; + uint32_t extra_meta_size_{0}; uint32_t element_size_{0}; }; From f08463011b0af38ecde1d987d9279a78f247aefe Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 15:46:13 +0800 Subject: [PATCH 05/45] refactor: update meta --- src/turbo/quantizer/quantizer.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index c6d1db2bc..572d372b5 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -24,11 +24,11 @@ #include #include "distance.h" -using namespace zvec::core; - namespace zvec { namespace turbo { +using namespace zvec::core; + //! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. constexpr uint32_t kQuantizerMagic = 0x52545A51u; //! Current quantizer serialization format version. From c3c4ca35c0d1c7389c9c96206a7527be97bf3596 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 16:22:51 +0800 Subject: [PATCH 06/45] refactor: add fp32 quantizer --- src/include/zvec/turbo/turbo.h | 15 +- .../fp32_quantizer/fp32_quantizer.cc | 161 ++++++++++++++++++ .../quantizer/fp32_quantizer/fp32_quantizer.h | 127 ++++++++++++++ src/turbo/turbo.cc | 22 ++- 4 files changed, 312 insertions(+), 13 deletions(-) create mode 100644 src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc create mode 100644 src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index df353bded..2eff572d5 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -71,15 +71,16 @@ enum class CpuArchType { }; DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type); + QuantizeType quantize_type, + CpuArchType cpu_arch_type = CpuArchType::kAuto); -BatchDistanceFunc get_batch_distance_func(MetricType metric_type, - DataType data_type, - QuantizeType quantize_type); +BatchDistanceFunc get_batch_distance_func( + MetricType metric_type, DataType data_type, QuantizeType quantize_type, + CpuArchType cpu_arch_type = CpuArchType::kAuto); -QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, - DataType data_type, - QuantizeType quantize_type); +QueryPreprocessFunc get_query_preprocess_func( + MetricType metric_type, DataType data_type, QuantizeType quantize_type, + CpuArchType cpu_arch_type = CpuArchType::kAuto); // Returns the SIMD kernel for the uniform quantizer on the current CPU for // the given output data_type, or nullptr if no SIMD implementation is diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc new file mode 100644 index 000000000..efd4e951e --- /dev/null +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -0,0 +1,161 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "quantizer/fp32_quantizer/fp32_quantizer.h" +#include +#include +#include +#include +#include +#include +#include "core/quantizer/record_quantizer.h" + +namespace zvec { +namespace turbo { + +int Fp32Quantizer::init(const IndexMeta &meta, + const ailego::Params & /*params*/) { + meta_ = meta; + + meta_.set_meta(IndexMeta::DataType::DT_FP32, meta.dimension()); + + auto metric_name = meta.metric_name(); + if (metric_name == "Cosine") { + extra_meta_size_ = EXTRA_META_SIZE_COSINE; + meta_.set_extra_meta_size(extra_meta_size_); + } + + // `meta.dimension()` may be either the raw dim (when the caller passes a + // bare meta, e.g. unit tests) or the inflated storage dim (data + extras) + // when an upstream converter such as CosineConverter has already appended + // a norm float and set meta.extra_meta_size(). Distinguish via the input + // meta's extra_meta_size: if it is already set, the dim is inflated and + // we strip the extras to recover the raw dim; otherwise the dim is raw. + if (meta.extra_meta_size() > 0) { + original_dim_ = meta.dimension() - meta.extra_meta_size() / sizeof(float); + } else { + original_dim_ = meta.dimension(); + } + + // Cache the distance dispatch for the new Quantizer interface. + dp_query_func_ = + get_distance_func(metric_from_name(metric_name), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + dp_query_batch_func_ = + get_batch_distance_func(metric_from_name(metric_name), DataType::kFp32, + QuantizeType::kDefault, CpuArchType::kAuto); + + return 0; +} + +int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (qmeta.unit_size() != sizeof(float)) { + return IndexError_Unsupported; + } + + // qmeta.dimension() may be the inflated (data + extras) dimension when the + // caller uses meta_.dimension() directly (e.g. HnswDistCalculator). Use the + // raw original dim we recorded at init() to avoid over-reading the query. + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + size_t byte_size = raw_dim * sizeof(float); + out->resize(byte_size); + std::memcpy(&(*out)[0], query, byte_size); + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_FP32, raw_dim, + static_cast(type_), extra_meta_size_); + + return 0; +} + +int Fp32Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + size_t byte_size = qmeta.dimension() * sizeof(float); + out->resize(byte_size); + std::memcpy(out->data(), in, byte_size); + return 0; +} + +DistanceImpl Fp32Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + auto metric = metric_from_name(meta_.metric_name()); + auto func = get_distance_func(metric, DataType::kFp32, QuantizeType::kDefault, + CpuArchType::kAuto); + if (!func) { + return DistanceImpl{}; + } + auto batch_func = get_batch_distance_func( + metric, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); + + // The query is assumed to be already quantized — copy it directly. + std::string quantized_query(static_cast(query), + qmeta.element_size()); + return DistanceImpl(std::move(func), std::move(batch_func), + std::move(quantized_query), original_dim_); +} + +void Fp32Quantizer::quantize_one(const void *input, void *output) const { + std::memcpy(output, input, + static_cast(original_dim_) * sizeof(float)); +} + +float Fp32Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + float d = 0.0f; + if (dp_query_func_) { + dp_query_func_(dp, query, original_dim_, &d); + } + return d; +} + +void Fp32Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const { + if (dp_query_batch_func_) { + dp_query_batch_func_(const_cast(dp_list), query, + static_cast(dp_num), original_dim_, dist_list); + return; + } + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = calc_distance_dp_query(dp_list[i], query); + } +} + +float Fp32Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + return calc_distance_dp_query(dp, buf.data()); +} + +void Fp32Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + calc_distance_dp_query_batch(dp_list, dp_num, buf.data(), dist_list); +} + +float Fp32Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + return calc_distance_dp_query(dp1, dp2); +} + +INDEX_FACTORY_REGISTER_QUANTIZER(Fp32Quantizer); + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h new file mode 100644 index 000000000..998a19d68 --- /dev/null +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -0,0 +1,127 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +class Fp32Quantizer : public Quantizer { + public: + Fp32Quantizer() { + type_ = QuantizeType::kFp32; + } + + virtual ~Fp32Quantizer() {} + + public: + int init(const core::IndexMeta &meta, const ailego::Params ¶ms) override; + + const core::IndexMeta &meta(void) const override { + return meta_; + } + + DataType input_data_type() const override { + return DataType::kFp32; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return false; + } + + int train(core::IndexHolder::Pointer /*holder*/) override { + return 0; + } + + size_t quantized_datapoint_vector_length() const override { + return quantized_length(); + } + + size_t quantized_query_vector_length() const override { + return quantized_length(); + } + + void quantize_data(const void *input, void *output) const override { + quantize_one(input, output); + } + + void quantize_query(const void *input, void *output) const override { + quantize_one(input, output); + } + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const core::IndexQueryMeta &qmeta, + std::string *out, core::IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const core::IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const core::IndexQueryMeta &qmeta) const override; + + private: + //! Byte length of a quantized vector (raw fp32 data). + size_t quantized_length() const { + return static_cast(original_dim_) * sizeof(float); + } + + //! Quantize a single fp32 vector into a caller-provided buffer of + //! quantized_length() bytes. + void quantize_one(const void *input, void *output) const; + + static constexpr uint32_t EXTRA_META_SIZE_COSINE = 4; + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + IndexMeta::DataType data_type_{}; + + //! Cached distance dispatch (bound in init()). + DistanceFunc dp_query_func_{}; + BatchDistanceFunc dp_query_batch_func_{}; +}; + + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index adc9b785e..11a7a43cc 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,10 +22,13 @@ namespace zvec::turbo { DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type) { + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_distance; } @@ -47,10 +50,13 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, BatchDistanceFunc get_batch_distance_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type) { + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_batch_distance; } @@ -67,15 +73,19 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, } } } + return nullptr; } QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type) { + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + (cpu_arch_type == CpuArchType::kAuto || + cpu_arch_type == CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_query_preprocess; } From ec4ea1fdfbfd6b5240cef161c1820b15f070c1f3 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 16:35:47 +0800 Subject: [PATCH 07/45] refactor: add quantizer --- src/turbo/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index fcf9dc431..7ec9c0e5d 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -29,7 +29,7 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH) endif() cc_library( - NAME zvec_turbo STATIC STRICT PACKED + NAME zvec_turbo STATIC STRICT ALWAYS_LINK PACKED SRCS ${ALL_SRCS} LIBS zvec_ailego INCS ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/distance ${PROJECT_ROOT_DIR}/src/include From e36f0c5495735fba1ba04e82149e9e39fcbac122 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 17:34:23 +0800 Subject: [PATCH 08/45] refactor: quantizer --- src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 998a19d68..79a7caa09 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -57,7 +57,7 @@ class Fp32Quantizer : public Quantizer { } int train(core::IndexHolder::Pointer /*holder*/) override { - return 0; + return IndexError_NotImplemented; } size_t quantized_datapoint_vector_length() const override { From c1cdd0b8c0fbbf337a85cc17bc4391cf234540e4 Mon Sep 17 00:00:00 2001 From: ray Date: Tue, 30 Jun 2026 19:08:56 +0800 Subject: [PATCH 09/45] refactor: add scalar --- src/turbo/distance/scalar/fp32/cosine.cc | 47 +++++++ src/turbo/distance/scalar/fp32/cosine.h | 30 +++++ .../distance/scalar/fp32/inner_product.cc | 43 +++++++ .../distance/scalar/fp32/inner_product.h | 31 +++++ .../distance/scalar/fp32/squared_euclidean.cc | 41 ++++++ .../distance/scalar/fp32/squared_euclidean.h | 31 +++++ src/turbo/turbo.cc | 33 +++++ tests/turbo/turbo_fp32_quantizer_test.cc | 119 ++++++++++++++++++ 8 files changed, 375 insertions(+) create mode 100644 src/turbo/distance/scalar/fp32/cosine.cc create mode 100644 src/turbo/distance/scalar/fp32/cosine.h create mode 100644 src/turbo/distance/scalar/fp32/inner_product.cc create mode 100644 src/turbo/distance/scalar/fp32/inner_product.h create mode 100644 src/turbo/distance/scalar/fp32/squared_euclidean.cc create mode 100644 src/turbo/distance/scalar/fp32/squared_euclidean.h diff --git a/src/turbo/distance/scalar/fp32/cosine.cc b/src/turbo/distance/scalar/fp32/cosine.cc new file mode 100644 index 000000000..112a6fc80 --- /dev/null +++ b/src/turbo/distance/scalar/fp32/cosine.cc @@ -0,0 +1,47 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp32/cosine.h" +#include + +namespace zvec::turbo::scalar { + +void cosine_fp32_distance(const void *a, const void *b, size_t dim, + float *distance) { + const float *fa = static_cast(a); + const float *fb = static_cast(b); + float dot = 0.0f; + float na = 0.0f; + float nb = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += fa[i] * fb[i]; + na += fa[i] * fa[i]; + nb += fb[i] * fb[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + if (denom < 1e-12f) { + *distance = 1.0f; + } else { + *distance = 1.0f - dot / denom; + } +} + +void cosine_fp32_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances) { + for (size_t i = 0; i < n; ++i) { + cosine_fp32_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/cosine.h b/src/turbo/distance/scalar/fp32/cosine.h new file mode 100644 index 000000000..b5e4f4eee --- /dev/null +++ b/src/turbo/distance/scalar/fp32/cosine.h @@ -0,0 +1,30 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute cosine distance (negative inner product after normalization) between +// a single quantized FP32 vector pair. +void cosine_fp32_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of cosine_fp32_distance. +void cosine_fp32_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances); + +} // namespace zvec::turbo::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/inner_product.cc b/src/turbo/distance/scalar/fp32/inner_product.cc new file mode 100644 index 000000000..63a0575a3 --- /dev/null +++ b/src/turbo/distance/scalar/fp32/inner_product.cc @@ -0,0 +1,43 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp32/inner_product.h" + +namespace zvec::turbo::scalar { + +// Compute squared Euclidean distance between a single quantized FP32 +// vector pair. +void inner_product_fp32_distance(const void *a, const void *b, size_t dim, + float *distance) { + const float *m = reinterpret_cast(a); + const float *q = reinterpret_cast(b); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i] * q[i]); + } + + *distance = -sum; +} + +// Batch version of inner_product_fp32_distance. +void inner_product_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, size_t dim, + float *distances) { + for (size_t i = 0; i < n; ++i) { + inner_product_fp32_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/inner_product.h b/src/turbo/distance/scalar/fp32/inner_product.h new file mode 100644 index 000000000..d4e03418e --- /dev/null +++ b/src/turbo/distance/scalar/fp32/inner_product.h @@ -0,0 +1,31 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute inner product distance between a single quantized FP32 +// vector pair. +void inner_product_fp32_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of inner_product_fp32_distance. +void inner_product_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, size_t dim, + float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp32/squared_euclidean.cc b/src/turbo/distance/scalar/fp32/squared_euclidean.cc new file mode 100644 index 000000000..602d8bd7f --- /dev/null +++ b/src/turbo/distance/scalar/fp32/squared_euclidean.cc @@ -0,0 +1,41 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp32/squared_euclidean.h" +#include + +namespace zvec::turbo::scalar { + +void squared_euclidean_fp32_distance(const void *a, const void *b, size_t dim, + float *distance) { + const float *m = reinterpret_cast(a); + const float *q = reinterpret_cast(b); + + float sum = 0.0; + for (size_t i = 0; i < dim; ++i) { + sum += zvec::ailego::MathHelper::SquaredDifference(m[i], q[i]); + } + + *distance = sum; +} + +void squared_euclidean_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, float *distances) { + for (size_t i = 0; i < n; ++i) { + squared_euclidean_fp32_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar \ No newline at end of file diff --git a/src/turbo/distance/scalar/fp32/squared_euclidean.h b/src/turbo/distance/scalar/fp32/squared_euclidean.h new file mode 100644 index 000000000..bf319c1d2 --- /dev/null +++ b/src/turbo/distance/scalar/fp32/squared_euclidean.h @@ -0,0 +1,31 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute squared euclidean distance between a single quantized FP32 +// vector pair. +void squared_euclidean_fp32_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of squared euclidean FP32. +void squared_euclidean_fp32_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 11a7a43cc..80f2ee77d 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -18,12 +18,30 @@ #include "avx512_vnni/record_quantized_int8/squared_euclidean.h" #include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/squared_euclidean.h" +#include "scalar/fp32/cosine.h" +#include "scalar/fp32/inner_product.h" +#include "scalar/fp32/squared_euclidean.h" namespace zvec::turbo { DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type) { + if (data_type == DataType::kFp32) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp32) { + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp32_distance; + } + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp32_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp32_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && @@ -52,6 +70,21 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type) { + if (data_type == DataType::kFp32) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp32) { + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp32_batch_distance; + } + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp32_batch_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp32_batch_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && diff --git a/tests/turbo/turbo_fp32_quantizer_test.cc b/tests/turbo/turbo_fp32_quantizer_test.cc index 40165a5d3..c180de4a4 100644 --- a/tests/turbo/turbo_fp32_quantizer_test.cc +++ b/tests/turbo/turbo_fp32_quantizer_test.cc @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include +#include #include #include #include @@ -22,6 +24,18 @@ using namespace zvec; using namespace zvec::core; using namespace zvec::ailego; +// Helper: reference cosine distance between two raw fp32 vectors. +static float reference_cosine(const float *a, const float *b, size_t dim) { + float dot = 0.0f, na = 0.0f, nb = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom; +} + TEST(Fp32Quantizer, General) { std::mt19937 gen(15583); std::uniform_real_distribution dist(0.0, 1.0); @@ -80,4 +94,109 @@ TEST(Fp32Quantizer, General) { EXPECT_NEAR(original_data[i], dequantize_data[i], 1e-3); } } +} + +TEST(Fp32Quantizer, Score) { + std::mt19937 gen(42); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Fp32Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0u, quantizer->init(meta, params)); + + // Generate raw vectors and quantize them. + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + IndexQueryMeta ometa; + EXPECT_EQ(0, quantizer->quantize( + raw_vecs[i].data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, DIMENSION), + &quant_vecs[i], &ometa)); + } + + // --- calc_distance_dp_query (single) --- + for (size_t i = 1; i < COUNT; ++i) { + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } + + // --- calc_distance_dp_query_batch --- + { + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch( + dp_list.data(), static_cast(dp_list.size()), quant_vecs[0].data(), + results.data()); + + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[i + 1].data(), + raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(results[i], expected, 1e-4) << "i=" << i; + } + } + + // --- distance() + DistanceImpl (single + batch) --- + { + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIMENSION); + auto dist_impl = quantizer->distance(quant_vecs[0].data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + + for (size_t i = 1; i < COUNT; ++i) { + float d = dist_impl(quant_vecs[i].data()); + float expected = + reference_cosine(raw_vecs[0].data(), raw_vecs[i].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } + + // Batch via DistanceImpl. + ASSERT_TRUE(dist_impl.batch_valid()); + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector batch_results(COUNT - 1); + dist_impl.batch(dp_list.data(), dp_list.size(), batch_results.data()); + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[0].data(), + raw_vecs[i + 1].data(), DIMENSION); + EXPECT_NEAR(batch_results[i], expected, 1e-4) << "i=" << i; + } + } + + // --- calc_distance_dp_dp (pairwise) --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_dp(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } + + // --- calc_distance_dp_query_unquantized --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_query_unquantized( + quant_vecs[i].data(), raw_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-4) << "i=" << i; + } } \ No newline at end of file From 628982db5738d6e431e37d362e9f85816d93e09c Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 11:53:21 +0800 Subject: [PATCH 10/45] fix: fix cosine --- src/turbo/distance/scalar/fp32/cosine.cc | 31 +++++++------------ .../quantizer/fp32_quantizer/fp32_quantizer.h | 2 +- 2 files changed, 13 insertions(+), 20 deletions(-) diff --git a/src/turbo/distance/scalar/fp32/cosine.cc b/src/turbo/distance/scalar/fp32/cosine.cc index 112a6fc80..6dd9fc532 100644 --- a/src/turbo/distance/scalar/fp32/cosine.cc +++ b/src/turbo/distance/scalar/fp32/cosine.cc @@ -13,34 +13,27 @@ // limitations under the License. #include "scalar/fp32/cosine.h" -#include +#include "scalar/fp32/inner_product.h" namespace zvec::turbo::scalar { void cosine_fp32_distance(const void *a, const void *b, size_t dim, float *distance) { - const float *fa = static_cast(a); - const float *fb = static_cast(b); - float dot = 0.0f; - float na = 0.0f; - float nb = 0.0f; - for (size_t i = 0; i < dim; ++i) { - dot += fa[i] * fb[i]; - na += fa[i] * fa[i]; - nb += fb[i] * fb[i]; - } - float denom = std::sqrt(na) * std::sqrt(nb); - if (denom < 1e-12f) { - *distance = 1.0f; - } else { - *distance = 1.0f - dot / denom; - } + // inner_product_fp32_distance returns -real_IP; cosine = 1 - real_IP = 1 + + // ip. + float ip; + inner_product_fp32_distance(a, b, dim, &ip); + + *distance = 1 + ip; } void cosine_fp32_batch_distance(const void *const *vectors, const void *query, size_t n, size_t dim, float *distances) { - for (size_t i = 0; i < n; ++i) { - cosine_fp32_distance(vectors[i], query, dim, &distances[i]); + inner_product_fp32_batch_distance(vectors, query, n, dim, distances); + // inner_product batch returns -real_IP per element; cosine = 1 - real_IP = 1 + // + d. + for (size_t i = 0; i < n; i++) { + distances[i] = 1 + distances[i]; } } diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 79a7caa09..998a19d68 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -57,7 +57,7 @@ class Fp32Quantizer : public Quantizer { } int train(core::IndexHolder::Pointer /*holder*/) override { - return IndexError_NotImplemented; + return 0; } size_t quantized_datapoint_vector_length() const override { From de3322042a2b7f300d99b6403c47666135676f91 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 14:17:05 +0800 Subject: [PATCH 11/45] fix: add quantizer --- src/turbo/distance/scalar/fp32/inner_product.cc | 4 ++-- src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc | 6 ------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/turbo/distance/scalar/fp32/inner_product.cc b/src/turbo/distance/scalar/fp32/inner_product.cc index 63a0575a3..1f966d788 100644 --- a/src/turbo/distance/scalar/fp32/inner_product.cc +++ b/src/turbo/distance/scalar/fp32/inner_product.cc @@ -16,8 +16,8 @@ namespace zvec::turbo::scalar { -// Compute squared Euclidean distance between a single quantized FP32 -// vector pair. +// Compute negated inner product between a single FP32 vector pair. +// Returns -dot(a, b) so that callers can derive cosine distance as 1 + ip. void inner_product_fp32_distance(const void *a, const void *b, size_t dim, float *distance) { const float *m = reinterpret_cast(a); diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index efd4e951e..6f8853c5c 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -36,12 +36,6 @@ int Fp32Quantizer::init(const IndexMeta &meta, meta_.set_extra_meta_size(extra_meta_size_); } - // `meta.dimension()` may be either the raw dim (when the caller passes a - // bare meta, e.g. unit tests) or the inflated storage dim (data + extras) - // when an upstream converter such as CosineConverter has already appended - // a norm float and set meta.extra_meta_size(). Distinguish via the input - // meta's extra_meta_size: if it is already set, the dim is inflated and - // we strip the extras to recover the raw dim; otherwise the dim is raw. if (meta.extra_meta_size() > 0) { original_dim_ = meta.dimension() - meta.extra_meta_size() / sizeof(float); } else { From 49780f897918d2006d0e9b1b1fd95eddf0801786 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 1 Jul 2026 14:59:26 +0800 Subject: [PATCH 12/45] fix: fix quantizer --- .../fp32_quantizer/fp32_quantizer.cc | 64 +++++++++++++++---- .../quantizer/fp32_quantizer/fp32_quantizer.h | 5 +- 2 files changed, 54 insertions(+), 15 deletions(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index 6f8853c5c..f5c29b76b 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -30,18 +31,14 @@ int Fp32Quantizer::init(const IndexMeta &meta, meta_.set_meta(IndexMeta::DataType::DT_FP32, meta.dimension()); + + original_dim_ = meta.dimension(); auto metric_name = meta.metric_name(); if (metric_name == "Cosine") { extra_meta_size_ = EXTRA_META_SIZE_COSINE; meta_.set_extra_meta_size(extra_meta_size_); } - if (meta.extra_meta_size() > 0) { - original_dim_ = meta.dimension() - meta.extra_meta_size() / sizeof(float); - } else { - original_dim_ = meta.dimension(); - } - // Cache the distance dispatch for the new Quantizer interface. dp_query_func_ = get_distance_func(metric_from_name(metric_name), DataType::kFp32, @@ -65,9 +62,22 @@ int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) ? original_dim_ : qmeta.dimension(); - size_t byte_size = raw_dim * sizeof(float); + size_t byte_size = raw_dim * sizeof(float) + extra_meta_size_; out->resize(byte_size); - std::memcpy(&(*out)[0], query, byte_size); + + if (meta_.metric_name() == "Cosine") { + // L2-normalize the vector in-place and store the norm at the end so the + // original vector can be reconstructed during dequantize. + float *buf = reinterpret_cast(&(*out)[0]); + std::memcpy(buf, query, raw_dim * sizeof(float)); + float norm = 0.0f; + ailego::Normalizer::L2(buf, raw_dim, &norm); + std::memcpy( + reinterpret_cast(&(*out)[0]) + raw_dim * sizeof(float), + &norm, extra_meta_size_); + } else { + std::memcpy(&(*out)[0], query, byte_size); + } *ometa = qmeta; ometa->set_meta(IndexMeta::DataType::DT_FP32, raw_dim, @@ -78,9 +88,26 @@ int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, int Fp32Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, std::string *out) const { - size_t byte_size = qmeta.dimension() * sizeof(float); - out->resize(byte_size); - std::memcpy(out->data(), in, byte_size); + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + size_t byte_size = raw_dim * sizeof(float); + + if (meta_.metric_name() == "Cosine") { + // Denormalize the vector using the stored norm. + out->resize(byte_size); + const float *in_buf = reinterpret_cast(in); + float *out_buf = reinterpret_cast(&(*out)[0]); + float norm = 0.0f; + std::memcpy(&norm, reinterpret_cast(in) + byte_size, + extra_meta_size_); + for (size_t i = 0; i < raw_dim; ++i) { + out_buf[i] = in_buf[i] * norm; + } + } else { + out->resize(byte_size); + std::memcpy(out->data(), in, byte_size); + } return 0; } @@ -103,8 +130,19 @@ DistanceImpl Fp32Quantizer::distance(const void *query, } void Fp32Quantizer::quantize_one(const void *input, void *output) const { - std::memcpy(output, input, - static_cast(original_dim_) * sizeof(float)); + size_t byte_size = static_cast(original_dim_) * sizeof(float); + + if (meta_.metric_name() == "Cosine") { + // L2-normalize and store the norm at the end. + std::memcpy(output, input, byte_size); + float *buf = reinterpret_cast(output); + float norm = 0.0f; + ailego::Normalizer::L2(buf, original_dim_, &norm); + std::memcpy(reinterpret_cast(output) + byte_size, &norm, + extra_meta_size_); + } else { + std::memcpy(output, input, byte_size); + } } float Fp32Quantizer::calc_distance_dp_query(const void *dp, diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 998a19d68..328a71e29 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -102,9 +102,10 @@ class Fp32Quantizer : public Quantizer { const core::IndexQueryMeta &qmeta) const override; private: - //! Byte length of a quantized vector (raw fp32 data). + //! Byte length of a quantized vector (raw fp32 data + extra meta). size_t quantized_length() const { - return static_cast(original_dim_) * sizeof(float); + return static_cast(original_dim_) * sizeof(float) + + extra_meta_size_; } //! Quantize a single fp32 vector into a caller-provided buffer of From b1d484956fa425aaba5c3c9726dd7046e4eeb32e Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 10:58:18 +0800 Subject: [PATCH 13/45] fix: remove unused var --- src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h | 1 - 1 file changed, 1 deletion(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 328a71e29..885390639 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -116,7 +116,6 @@ class Fp32Quantizer : public Quantizer { IndexMeta meta_{}; uint32_t original_dim_{0}; - IndexMeta::DataType data_type_{}; //! Cached distance dispatch (bound in init()). DistanceFunc dp_query_func_{}; From f0fdf5d65fe8663ed89b91b43d8b25635b346e25 Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 14:08:42 +0800 Subject: [PATCH 14/45] fix: header --- examples/c++/CMakeLists.txt | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index d58814d52..4de64b45c 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -14,7 +14,10 @@ if(NOT DEFINED HOST_BUILD_DIR) endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) -set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) +# Include both src/include (public headers) and src (internal headers that +# are transitively reachable from public headers, e.g. +# turbo/quantizer/quantizer.h via index_factory.h). +set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include ${ZVEC_ROOT_DIR}/src) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) # Add include and library search paths From 0c4cec659b4e285b668411ab829f1bc00e07311e Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 15:05:02 +0800 Subject: [PATCH 15/45] fix: header --- src/turbo/quantizer/quantizer.h | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 572d372b5..14e889787 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -18,7 +18,6 @@ #include #include #include -#include #include #include #include @@ -31,6 +30,7 @@ using namespace zvec::core; //! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. constexpr uint32_t kQuantizerMagic = 0x52545A51u; + //! Current quantizer serialization format version. constexpr uint16_t kQuantizerSerVersion = 1; @@ -75,7 +75,7 @@ class Quantizer { //! Train the quantizer with a contiguous batch of data virtual int train(const void * /*data*/, size_t /*num*/, size_t /*stride*/) { - return IndexError_NotImplemented; + return 0; } //! Whether the quantizer requires training before use @@ -83,7 +83,7 @@ class Quantizer { //! Train the quantizer with data from an IndexHolder virtual int train(IndexHolder::Pointer /*holder*/) { - return IndexError_NotImplemented; + return 0; } //! Byte length of a quantized datapoint vector @@ -123,13 +123,13 @@ class Quantizer { virtual int quantize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, std::string * /*out*/, IndexQueryMeta * /*ometa*/) const { - return IndexError_NotImplemented; + return 0; } //! Dequantize a result vector back to original format virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, std::string * /*out*/) const { - return IndexError_NotImplemented; + return 0; } virtual DistanceImpl distance(const void * /*query*/, @@ -139,18 +139,18 @@ class Quantizer { //! Serialize quantizer parameters virtual int serialize(std::string * /*out*/) const { - return IndexError_NotImplemented; + return 0 } //! Deserialize quantizer parameters virtual int deserialize(std::string & /*in*/) { - return IndexError_NotImplemented; + return 0; } //! Deserialize quantizer parameters from a raw, possibly mmap-backed buffer //! (zero-copy entry point for large payloads such as codebooks/matrices). virtual int deserialize(const void * /*data*/, size_t /*len*/) { - return IndexError_NotImplemented; + return 0; } protected: From 73f5bc3bab498315bac70d8c3142a7ecb530020c Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 15:36:15 +0800 Subject: [PATCH 16/45] fix: header --- src/turbo/quantizer/quantizer.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 14e889787..3e5e65ec6 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -139,7 +139,7 @@ class Quantizer { //! Serialize quantizer parameters virtual int serialize(std::string * /*out*/) const { - return 0 + return 0; } //! Deserialize quantizer parameters From 89843dd92a2131921ad2bcfdefaf93d73e827010 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 2 Jul 2026 15:42:47 +0800 Subject: [PATCH 17/45] add preprocessor --- ...26\345\231\250\350\256\276\350\256\241.md" | 731 ++++++++++++++++++ doc/pq.md | 94 +++ src/include/zvec/turbo/turbo.h | 20 + src/turbo/CMakeLists.txt | 23 +- src/turbo/distance/avx2/fht/fht.h | 37 + src/turbo/distance/avx2/fht/fht_avx2.cc | 138 ++++ src/turbo/distance/avx512/fht/fht.h | 37 + src/turbo/distance/avx512/fht/fht_avx512.cc | 143 ++++ src/turbo/distance/scalar/fht/fht.h | 38 + src/turbo/distance/scalar/fht/fht_scalar.cc | 80 ++ src/turbo/distance/sse/fht/fht.h | 34 + src/turbo/distance/sse/fht/fht_sse.cc | 119 +++ .../quantizer/preprocessor/fht_rotator.cc | 266 +++++++ .../quantizer/preprocessor/fht_rotator.h | 97 +++ .../quantizer/preprocessor/preprocessor.h | 93 +++ src/turbo/turbo.cc | 54 ++ tests/turbo/turbo_fht_rotator_test.cc | 204 +++++ 17 files changed, 2205 insertions(+), 3 deletions(-) create mode 100644 "doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" create mode 100644 doc/pq.md create mode 100644 src/turbo/distance/avx2/fht/fht.h create mode 100644 src/turbo/distance/avx2/fht/fht_avx2.cc create mode 100644 src/turbo/distance/avx512/fht/fht.h create mode 100644 src/turbo/distance/avx512/fht/fht_avx512.cc create mode 100644 src/turbo/distance/scalar/fht/fht.h create mode 100644 src/turbo/distance/scalar/fht/fht_scalar.cc create mode 100644 src/turbo/distance/sse/fht/fht.h create mode 100644 src/turbo/distance/sse/fht/fht_sse.cc create mode 100644 src/turbo/quantizer/preprocessor/fht_rotator.cc create mode 100644 src/turbo/quantizer/preprocessor/fht_rotator.h create mode 100644 src/turbo/quantizer/preprocessor/preprocessor.h create mode 100644 tests/turbo/turbo_fht_rotator_test.cc diff --git "a/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" "b/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" new file mode 100644 index 000000000..66f02415e --- /dev/null +++ "b/doc/Zvec Turbo \351\207\217\345\214\226\345\231\250\350\256\276\350\256\241.md" @@ -0,0 +1,731 @@ +# 背景 +当前，距离度量沿用了 Proxima 的风格，采用模板方式进行定义。系统通过模块化设计实现了度量、转换器及指令集等组件的解耦，检索过程由Metric（距离度量）、Quantizer（量化器）、Algorithm（索引算法）等多个模块协作完成。然而，一些现有的组件存在职责边界模糊、实现方式耦合等问题，计划通过改为turbo的形式，统一切面通过输入多个维度信息组合以获得对应的函数集合，从而使得不同逻辑判断统一。 + +# 设计现状 +当前的系统架构设计存在着一些问题，下面将进行简要说明。 + +## Converter和Reformer功能耦合 +当前的量化器由Converter和Reformer两个组件构成，其中，Converter主要负责离线构建，而Reformer主要负责在线查询，不过由于当前设计未区分流式插入与批量合并，Converter和Reformer在离线阶段存在一定的功能耦合： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/226557004/1782716498906-f0957db9-24f8-466b-bfb2-351e65d1b5c5.jpeg) + +这个实现最初的想法是将离线的一次性操作`build`、`train`等放到Converter中，而可以反复进行的操作`add`和`search`等放到Reformer，但是，这在本身也引入了一些额外的限制，例如，如果量化器需要训练，但索引是流式插入的，例如hnsw+kUniformInt8，就会存在问题。以下对当前的使用接口进行说明： + +| **Converter接口** | **功能** | **Reformer接口** | **功能** | +| --- | --- | --- | --- | +| init() | 声明量化后的 meta 信息 | init() | 从 Converter 声明的 reformer_params 中加载运行时参数 | +| train() | 在全量数据上执行训练，收集统计信息 | load() | 从持久化存储中加载/卸载 Reformer 运行时状态 | +| transform() | 将 全量向量从原始格式转换为量化格式 | transform() | 将查询向量转换为与索引一致的量化格式 | +| dump() | 将 Converter 内部状态（如量化参数）持久化到 dumper | normalize() | 将量化空间计算的距离/分数映射回原始距离尺度 | +| meta() | 返回 Converter 修改后的 IndexMeta | revert() | 将量化向量反量化回 FP32 | + + +## Quantizer与Metric存在耦合 +当前设计存在不同功能分布在同一文件/同一功能分布在不同文件的问题，例如： + +| 注册文件 | 工厂注册名 | Metric | 备注 | +| --- | --- | --- | --- | +| **integer_quantizer_converter.cc** | Int8StreamingConverter | L2 / IP / MIPS | 基于单条数据的INT8 | +| | Int4StreamingConverter | L2 / IP / MIPS | 基于单条数据的INT4 | +| | Int8QuantizerConverter | - | 全局数据的INT8 | +| | Int4QuantizerConverter | - | 全局数据的INT4 | +| **cosine_converter.cc** | CosineInt8Converter | Cosine | 基于单条数据的INT8 | +| | CosineInt4Converter | Cosine | 基于单条数据的INT4 | +| | CosineFp16Converter | Cosine | FP16 | +| | ... | ... | ... | + + +由于这样的设计，一些量化器内部需要判断度量类型、数据类型、架构类型，以合理分配量化类型。 + +## 其他问题 +1. Builder、Searcher、Streamer职责边界模糊： + 1. Builder：全量数据的构建，负责操作 `build` + 2. Searcher：只被 IndexFlow（CLI 工具）和单元测试使用 + 3. Streamer：主路径组件负责读写，负责操作 `add` 和`search` +2. 部分索引实现与量化器耦合： + 1. 与Quantizer完全解耦：HNSW / Flat / Vamana + 2. 硬编码与Quantizer耦合：IVF + 3. 当前自成一体：diskann/hnsw rabitq + +# Turbo设计 +Turbo 旨在通过统一接口收敛度量（Metric）与量化（Quantization）逻辑，提高代码内聚性，从而降低维护成本，提升代码的可读性与简洁性。 + +## 使用场景 +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782788752754-8c017988-c7ed-469d-8025-b1ab5ce94281.jpeg) + +## 模块 +Turbo系统包含以下几个模块： + ++ 量化器 ++ 预处理器 ++ 距离计算 + +都以量化器为核心进行调度： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782802278889-9e04f816-ec15-42f6-9ff6-3e53ceb7c83e.jpeg) + +## 设计目标 +设计的目标主要包含以下六个维度的： + +| **维度** | **说明** | +| :---: | :---: | +| 源数据类型 | 量化前的数据类型 | +| 目标数据类型 | 量化后的数据类型 | +| 距离类型 | 支持的距离度量方式 | +| 量化器类型 | 具体量化算法及参数管理方式 | +| 距离计算模式 | 单查询单点、单查询多点连续、单查询多点离散、点对点 | +| 指令架构 | 该量化器支持的指令架构 | +| 量化器使用场景 | 训练、流式插入、批量插入/合并、序列化、反序列化 | + + +## 调用场景 +以下为构建、查询的调用场景： + +### 构建 +```plain + 源向量 + │ + ▼ +┌───────────────────────┐ +│ Quantizer │ +│ .quantize_datapoint │ +└───────────────────────┘ + │ + ▼ +[量化向量 + Extra Meta(如有)] → 写入 Streamer +``` + +要求： + ++ Uniform：无需训练即可编码（全局参数已确定）。 ++ Record：逐向量独立计算 scale/bias，无需全局训练。 ++ PQ/RabitQ：不支持纯流式插入（需先训练）。 + +### 查询 +```plain + 源向量 + │ + ▼ +┌─────────────────────┐ +│ Quantizer │ +│ .quantize_query() │ +└─────────────────────┘ + │ + ▼ +[量化向量 + Extra Meta(如有)] → 检索 Streamer +``` + +### 合并 +```plain +多个源 Segment (FP32) + │ + ▼ +┌─────────────┐ +│ write_holder│ ← 聚合数据到Holder +└─────────────┘ + │ + ▼ +┌─────────────┐ +│ Quantizer │ ← 训练全局参数 +│ .train │ +└─────────────┘ + │ + ▼ +┌───────────────────────┐ +│ Quantizer │ ← 用训练后的Quantizer编码所有向量 +│ .quantize_datapoint | +└───────────────────────┘ + │ + ▼ +[量化向量 + Extra Meta(如有)] → 写入 Streamer +``` + +# 支持范围 +## 源数据类型 → 目标数据类型 +考虑到 db 层有不保留 Flat FP32 的需求 + ++ 源数据类型需要支持 4 种通用类型： + - FP32，P0 + - FP16，P0 + - INT8，P1 + - INT4，P2 ++ 目标数据类型支持范围： + +| **目标类型** | **说明** | **优先级** | +| :---: | :---: | :---: | +| FP32 | 原始数据不做量化，仅做度量转换（如 Cosine） | P0 | +| FP16 | 半精度浮点 | P0 | +| INT8 | 8-bit 整数量化 | P0 | +| INT4 | 4-bit 整数量化 | P0 | +| X-Bits | 自定义 bit 宽度的量化（如 5-bit、6-bit） | P1 | + + +## 量化器类型 +量化器类型是核心分类维度，决定了端到端的数据布局、训练方式和距离计算方式。 + +| **量化器类型** | **特征** | **是否带Extra Meta** | **距离计算方式** | **优先级** | +| :---: | :---: | :---: | :---: | :---: | +| **Record** | 逐向量保存 scale/bias，按维度量化 | 是
（Extra Meta 存 scale/bias） | 非对称距离，需反量化或查表 | P0 | +| **Uniform** | 全局统一scale & bias，按维度量化 | 否 | 对称距离，可 SIMD 优化 | P0 | +| **PQ**
** (Product Quantization)** | 分段聚类码本量化 | 是
（码本参数） | 非对称距离，ADC/SDC | P1 | +| **OPQ** | 优化的 PQ，含正交变换 | 是
（变换矩阵 + 码本） | 非对称距离，ADC | P2 | +| **Raw** | 不做数值量化，仅改变类型（FP32→FP16）或做度量预处理 | 否 | 原始距离计算 | P0 | +| ~~**RabitQ**~~ | ~~旋转 + 二进制量化~~ | ~~是（旋转矩阵等）~~ | ~~汉明距离近似~~ | ~~P1~~ | + + +_**注意**__：量化器类型与目标数据类型是正交但有关联的概念。例如 Record 和 Uniform 都可以输出到 INT8 或 INT4；Raw 可以输出到 FP32、FP16、BF16。_ + +## 距离类型支持 +| 距离类型 | 支持说明 | 优先级 | +| --- | --- | --- | +| **L2 (Squared Euclidean)** | 所有量化方案的基础支持 | P0 | +| **Cosine** | 所有量化方案支持，可通过预处理归一化转化为 IP 或 L2 | P0 | +| **Inner Product (IP)** | 所有量化方案支持 | P0 | +| **MIPS** | 基于 IP 的变形，需做转换 | P1 | + + +# 量化器 +## 接口语义 +每个具体的量化器实现以下统一接口，但内部行为因方案而异： + +```cpp +class Quantizer { + public: + typedef std::shared_ptr Pointer; + + Quantizer() {} + virtual ~Quantizer() {} + + //! Initialize quantizer with index metadata and parameters + virtual int init(const IndexMeta &meta, const ailego::Params ¶ms) = 0; + + //! Get the output metadata after initialization + virtual const IndexMeta &meta() const = 0; + + //! Input data type accepted by the quantizer + virtual DataType input_data_type() const = 0; + + //! Data type + virtual QuantizeType type() const { + return type_; + } + + //! Dimensionality of the input vectors + virtual int dim() const = 0; + + //! Train the quantizer with a contiguous batch of data + virtual int train(const void * /*data*/, size_t /*num*/, size_t /*stride*/) { + return IndexError_NotImplemented; + } + + //! Whether the quantizer requires training before use + virtual bool require_train() const = 0; + + //! Train the quantizer with data from an IndexHolder + virtual int train(IndexHolder::Pointer /*holder*/) { + return IndexError_NotImplemented; + } + + //! Byte length of a quantized datapoint vector + virtual size_t quantized_datapoint_vector_length() const = 0; + + //! Byte length of a quantized query vector + virtual size_t quantized_query_vector_length() const = 0; + + //! Quantize a datapoint vector + virtual void quantize_data(const void *input, void *output) const = 0; + + //! Quantize a query vector + virtual void quantize_query(const void *input, void *output) const = 0; + + //! Distance between a quantized datapoint and a quantized query + virtual float calc_distance_dp_query(const void *dp, + const void *query) const = 0; + + //! Batched distance between quantized datapoints and a quantized query + virtual void calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const = 0; + + //! Distance between a quantized datapoint and an unquantized query + virtual float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const = 0; + + //! Batched distance between quantized datapoints and an unquantized query + virtual void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const = 0; + + //! Distance between two quantized datapoints + virtual float calc_distance_dp_dp(const void *dp1, const void *dp2) const = 0; + + //! Quantize a query vector for search + virtual int quantize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/, + IndexQueryMeta * /*ometa*/) const { + return IndexError_NotImplemented; + } + + //! Dequantize a result vector back to original format + virtual int dequantize(const void * /*in*/, const IndexQueryMeta & /*qmeta*/, + std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + virtual DistanceImpl distance(const void * /*query*/, + const IndexQueryMeta & /*qmeta*/) const { + return DistanceImpl{}; + } + + //! Serialize quantizer parameters + virtual int serialize(std::string * /*out*/) const { + return IndexError_NotImplemented; + } + + //! Deserialize quantizer parameters + virtual int deserialize(std::string & /*in*/) { + return IndexError_NotImplemented; + } + + //! Deserialize quantizer parameters from a raw, possibly mmap-backed buffer + //! (zero-copy entry point for large payloads such as codebooks/matrices). + virtual int deserialize(const void * /*data*/, size_t /*len*/) { + return IndexError_NotImplemented; + } + + protected: + //! Map a metric name (e.g. "SquaredEuclidean", "Cosine", + //! "InnerProduct", "MipsSquaredEuclidean") to its MetricType. + static MetricType metric_from_name(const std::string &name) { + if (name == "SquaredEuclidean") { + return MetricType::kSquaredEuclidean; + } + if (name == "Cosine") { + return MetricType::kCosine; + } + if (name == "InnerProduct") { + return MetricType::kInnerProduct; + } + if (name == "MipsSquaredEuclidean") { + return MetricType::kMipsSquaredEuclidean; + } + return MetricType::kUnknown; + } + + QuantizeType type_{QuantizeType::kDefault}; + uint32_t extra_meta_size_{0}; +}; +``` + +dp * dp batch for IVF -> P1 + +## 量化方案 +#### Raw 方案（类型转换/度量预处理） ++ **Extra Meta**：对于Cosine向量增加归一化值。 ++ **编码流程**： + - FP32 → FP16：数值截断/舍入。 + - FP32 → FP32 + Cosine 预处理：向量归一化。 ++ **距离计算**：使用对应精度的原始距离计算（如 FP16 L2、Cosine 转 IP）。 ++ **使用场景**：无需训练，实时转换。不需要训练，支持流式 + +#### Record 方案（逐向量 scale/bias） ++ **Extra Meta**：每个向量后附加 scale、bias相关的数组（或其他压缩形式）。 ++ **编码流程**：输入 FP32 → 计算向量的scale/bias参数 → 量化到 INT8/INT4 → 输出 [codes + scale + bias]。 ++ **距离计算**：非对称，查询向量需经过相同 scale/bias 反量化或预处理到同一域。 ++ **使用场景**：流式插入友好（每条向量独立计算 meta），但 Extra Meta 增加存储。 + +#### Uniform 方案（全局 scale/bias） ++ **Extra Meta**：无 per-vector meta，全局参数存于 Quantizer 内部（通过 serialize/deserialize 持久化）。 ++ **编码流程**：输入 FP32 → 使用全局 scale/bias → 量化到 INT8/INT4 → 纯码本输出。 ++ **距离计算**：对称距离，查询向量和库向量使用同一组全局参数，可高度 SIMD 优化。 ++ **使用场景**：批量训练确定全局参数后，流式插入和检索效率最高。 + +#### PQ / OPQ 方案 ++ **Extra Meta**：无 per-vector meta，全局码本参数（centroids）存于 Quantizer 内部（通过 serialize/deserialize 持久化）。 ++ **预处理：**OPQ有随机旋转预处理 ++ **编码流程**：输入 FP32 → （OPQ 先做正交变换）→ 分段聚类 → 输出码本索引。 ++ **距离计算**：ADC（Asymmetric Distance Computation）~~或 SDC~~。 ++ **使用场景**：高维向量压缩，适合批量构建，流式插入需支持码本增量更新（P2）。 ++ **实现计划**：[PQ in Zvec Turbo](https://yuque.alibaba-inc.com/proxima/aby36h/dt2x3o15c1mbqgxz) + +#### TurboQuant方案 ++ **Extra Meta：**无per-vector meta， ++ **预处理：**需要加随机旋转 ++ **距离计算：**SIMD， INT计算 ++ **使用场景：**极致压缩，保持一定召回。不需要训练，支持流式 + +#### ~~RabitQ 方案~~ ++ **Extra Meta**：旋转矩阵等参数。 ++ **编码流程**：输入 FP32 → 旋转 → 二进制量化。 ++ **距离计算**：基于汉明距离的近似计算。 ++ **使用场景**：极高压缩比，精度要求可容忍的场景。 + + + +## 训练 +量化器训练过程会在合并阶段被调用。量化器根据各自实现，来确定是否需要训练。量化器的训练过程，通过传入一组连续内存的数据，并用步长来指定采样方式，进行训练。 + +# 距离计算器 +为了对距离方式进行封装，设计了距离计算器DistanceImpl类： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/95883/1782875740300-e562d52a-5bfc-485d-b6cf-61e523e3edaf.jpeg) + +# 距离表 +基于上述分类，Turbo 的函数分发不再通过分散的模板特化，而是通过**量化方案 + 目标数据类型 + 距离类型 + SIMD 架构**四个维度组合查询函数表： + +类型定义： + ++ 度量类型MetricType定义： + - kSquaredEuclidean + - kCosine + - kInnerProduct + - kMipsSquaredEuclidean + - kUnknown ++ 数据类型DataType定义 + - kInt4 + - kInt8 + - kFp16 + - kFp32 + - kUnknown ++ 量化类型QuantizeType定义 + - kDefault + - kUniform + - kRecord + - kFp16 + - kFp32 + - kPQ + - kRabit ++ 架构类型CpuArchType定义： + - kAuto + - kScalar + - kSSE + - kAVX + - kAVX2 + - kAVX512 + - kAVX512VNNI + - kAVX512FP16 + +相应的，获取函数的接口也需要枚举对应的参数类型： + ++ DistanceFunc **get_distance_func**(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); ++ BatchDistanceFunc **get_batch_distance_func**(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); ++ QueryPreprocessFunc **get_query_preprocess_func**(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); + +**说明**： + ++ 不存在的组合在编译期报错（static_assert），避免运行时意外。 ++ 运行时通过 Quantizer 返回的 DistanceCalculator 绑定已确定的 scheme/dst_type/metric，仅需March进行动态分发。 + +# 原值构建支持 +由于使用量化值进行图索引构建，会出现不够精确的情况，进一步支持使用原值而非量化值来进行。streamer 和 builder 在实现时严格区分 rdp * rdp （fp32 raw quantizer）; query*dp ，并提供设置原值索引 provider 的接口（之后可考虑和 external_storage 需求中设计的VectorSource 合并） + +术语：query / dp / r(eference)dp + +# 预处理器 +区分概念：预处理操作和预处理器操作，预处理操作包括：旋转、normalize，数据类型偏移等，但并不是所有的预处理操作都需要预处理器来完成，只有那些有状态依赖的操作（例如：旋转）需要使用旋转完成，而normalize，数据类型偏移等操作使用函数指针来完成更加便捷。 + +预处理器是一种量化器的内部可选组件，主要负责在进行数据检索前对数据进行预处理，主要包括旋转功能，可以对量化前的数据进行预处理，以帮助量化器发挥更好的效果。其中，在需要进行预处理时，量化器**主动**调用预处理器，在进行量化前对向量数据进行处理；需要保存时，预处理器需要的数据同量化器中的参数一起进行保存。该组件需要实现的功能包括： + ++ 随机旋转（optional）+ Record INT8/INT4 ++ OPQ旋转（optional）+ PQ ++ 通过某种方式，对数据进行降维（optional） + +职责边界：哪些功能属于预处理范畴，与索引/量化的边界在何处？预处理器是一种**可选**的量化方法的优化组件，缺少也不影响量化方法的正常运行，由量化方案创建的组件属于量化器。 + +预处理器中，处理数据同样是一个费时的操作，尤其是在在线处理的情况下，需要使用**距离表**中的方案进行加速，在获取相应的预处理函数的接口时，也需要枚举对应的参数类型，比如： + ++ FwhtFunc get_fwht_func(CpuArchType cpu_arch_type = CpuArchType::kAuto); ++ MatVecFunc get_matvec_func(CpuArchType cpu_arch_type = CpuArchType::kAuto); + +通过分发的方式可以快速获得对应的预处理函数。具体的实现如下： + +![画板](https://intranetproxy.alipay.com/skylark/lark/0/2026/jpeg/226557004/1782812444145-e4b10b66-8bd1-4f6d-aa94-c205ca57235d.jpeg) + +由于当前FHT可以适用于任意维度的向量，不必再将FHT和Matrix作为rotator下的子类，去除rotator，将FHT和Matrix作为两个独立的类。对于INT8/INT4使用FHT，而对于OPQ使用Matrix，未来还可能有降维等功能。 + +# RabitQ量化接入挑战 +针对当前HnswRabitq索引接入RabitQ量化器，主要有两个方面的挑战： + +1. 距离计算是渐进进行，先计算粗略距离，再进行精确距离计算，这个涉及到在索引遍历时的判断； +2. 返回结果时距离组合，对于流程中的距离返回需要进行改造，涉及到TopK队列的结构体定义； + +结论：暂时在第一期不做接入。 + +# IndexMeta调整 +为支持上述设计，IndexMeta 需增加以下字段： + +| 字段 | 类型 | 说明 | +| --- | --- | --- | +| quantize_type | QuantizeType | 量化方案（Record/Uniform/PQ...） | +| extra_meta_size | uint32_t | 每条向量的 Extra Meta 大小（0 表示无） | + + +兼容性处理： + ++ 存量数据无 quantize_type 字段时，按旧逻辑推断（如根据数据类型和是否有 Extra Meta 推断为 Record 或 Raw）。 ++ 新增字段以扩展字段方式存入 IndexMeta，保证向后兼容。 + +# 代码说明 +## 距离计算器定义 +```cpp +class DistanceImpl { + public: + DistanceImpl() = default; + + DistanceImpl(DistanceFunc func, std::string quantized_query, size_t dim) + : func_(std::move(func)), + query_storage_(std::move(quantized_query)), + dim_(dim) {} + + DistanceImpl(DistanceFunc func, BatchDistanceFunc batch_func, + std::string quantized_query, size_t dim) + : func_(std::move(func)), + batch_func_(std::move(batch_func)), + query_storage_(std::move(quantized_query)), + dim_(dim) {} + + //! Whether the handle is ready to compute distances. + bool valid() const { + return static_cast(func_); + } + + //! Whether a batch distance function is available. + bool batch_valid() const { + return static_cast(batch_func_); + } + + //! Compute the distance between the stored query and `candidate`. + float operator()(const void *candidate) const { + float d = 0.0f; + func_(candidate, query_storage_.data(), dim_, &d); + return d; + } + + //! Compute distances for a batch of `num` candidates against the + //! stored query. Falls back to the scalar path when no batch function + //! is bound. + void batch(const void **candidates, size_t num, float *out) const { + if (batch_func_) { + batch_func_(candidates, query_storage_.data(), num, dim_, out); + return; + } + for (size_t i = 0; i < num; ++i) { + out[i] = 0.0f; + func_(candidates[i], query_storage_.data(), dim_, out + i); + } + } + + //! Access the quantized query bytes (for pairwise helpers). + const std::string &query_storage() const { + return query_storage_; + } + + size_t dim() const { + return dim_; + } + + //! Raw scalar distance function (operates on already-quantized + //! candidates). Useful for pairwise node-vs-node distance where no + //! stored query is involved. + const DistanceFunc &func() const { + return func_; + } + + //! Raw batch distance function. + const BatchDistanceFunc &batch_func() const { + return batch_func_; + } + + private: + DistanceFunc func_{}; + BatchDistanceFunc batch_func_{}; + std::string query_storage_{}; + size_t dim_{0}; +}; +``` + +## 索引对接示例 +在索引中，通过接入量化器来实现距离对象的接入： + +```cpp +class HnswDistCalculator { + ... .. + private: + zvec::turbo::DistanceImpl dist_impl_{}; +}; +``` + +初始化： + +```cpp +inline void reset_query(const void *query) { + if (quantizer_) { + dist_impl_ = quantizer_->distance(query, qmeta_); + } + ... ... +} +``` + + 距离计算调用： + +```cpp +const auto &func = dist_impl_.func(); +if (func) { + func(vec_lhs, vec_rhs, dist_impl_.dim(), &score); + return score; +} +``` + +## 量化器实现示例 +对于有状态的距离计算，可以在量化器中，通过lambda来捕获对应的状态数据。 + +```cpp + +DistanceImpl MyQuantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + // Create a lambda-based DistanceFunc that captures this quantizer + auto dim = padded_dim_; + auto metric = metric_type_; + auto cb_val = cb(); + + DistanceFunc func = [dim, metric, cb_val](const void *m, const void *q, + size_t /*dim_arg*/, float *out) { + const auto *dp_bytes = reinterpret_cast(m); + const auto *q_floats = reinterpret_cast(q); + + const auto *code = reinterpret_cast(dp_bytes); + const float *factors = reinterpret_cast(dp_bytes + dim); + float f_add = factors[0]; + float f_rescale = factors[1]; + + const float *rotated_q = q_floats; + float sum_q = q_floats[dim]; + float norm_sq_q = q_floats[dim + 1]; + + float ip = ip_float_code(rotated_q, code, dim); + float est_dist = f_add + f_rescale * (ip + cb_val * sum_q); + if (metric == MetricType::kSquaredEuclidean) { + est_dist += norm_sq_q; + } + *out = est_dist; + }; + + return DistanceImpl(std::move(func), std::move(qbuf), padded_dim_); +} +``` + +## 预处理示例 +增加预处理来支持向量的前置处理过程，其中旋转也是预处理的一部分。定义如下： + +```cpp +class Preprocessor { + public: + using Pointer = std::shared_ptr; + + virtual ~Preprocessor() = default; + + //! Input dimensionality accepted by apply(). + virtual int in_dim() const = 0; + + //! Output dimensionality produced by apply(). May differ from in_dim() + //! (e.g. FHT pads to the next power of two). + virtual int out_dim() const = 0; + + //! Forward transform: map an input vector to the preprocessed space. + //! \p out must hold at least out_dim() elements. + virtual void apply(const float *in, float *out) const = 0; + + //! Inverse transform: recover the original-space vector from a preprocessed + //! one. \p out must hold at least in_dim() elements. + virtual void apply_inverse(const float *in, float *out) const = 0; + + //! Fit the preprocessor from a contiguous batch of training data. + //! \p data pointer to the first element of the batch. + //! \p num number of vectors in the batch. + //! \p stride byte offset between consecutive vectors (0 ⇒ packed). + virtual void train(const void *data, size_t num, size_t stride) = 0; + + //! Serialize the preprocessor into a self-contained blob (header + payload). + virtual int serialize(std::string *out) const = 0; + + //! Deserialize the preprocessor from a raw, possibly mmap-backed buffer. + virtual int deserialize(const void *data, size_t len) = 0; +}; +``` + +对于旋转定义： + +```cpp +class Rotator : public Preprocessor { + public: + using Pointer = std::shared_ptr; + + //! Kind of rotator. + virtual RotatorType type() const = 0; + + // in_dim(), out_dim(), apply(), apply_inverse(), train(), serialize(), + // and deserialize() are inherited from Preprocessor. +}; + +//! Create an untrained rotator of the given type for in_dim dimensions. +Rotator::Pointer CreateRotator(RotatorType type, int in_dim); + +//! Create and restore a rotator from a serialized blob (reads the type from +//! the embedded RotatorSerHeader). Returns nullptr on malformed input. +Rotator::Pointer CreateRotatorFromBlob(const void *data, size_t len); + +} // namespace turbo +``` + +预处理器作为量化器的成员，进行接入。 + +# 实现优先级 +## 一期（P0） ++ **量化方案**： + - Raw（FP32→FP16, FP32→FP32 Cosine） + - Uniform（FP32→INT8, FP32→INT4） + - Record（FP32→INT8, FP32→INT4） ++ **距离类型**：L2, Cosine, IP ++ **微架构类型：**SSE，AVX，AVX512，AVX512FP16 ++ **距离计算模式**：所有单点查询、部分批量、 ++ **索引支持：**HNSW，IVF，Vamana，DiskAnn ++ **使用场景**：流式插入、Merge 训练、流式检索、序列化/反序列化 + +## 二期（P1） ++ **量化方案**： + - PQ（FP32→PQ） ++ **距离类型**：L2, Cosine, IP ++ **微架构类型：**SSE，AVX，AVX512，AVX512FP16 ++ **距离计算模式**：所有批量查询 + +# 模块分工 +在基础PR的基础上，其它模块可以分PR进行提交、验证和合并。 + +| **模块** | **内容** | **人员** | **开始日期** | **结束日期** | **人日** | **优先级** | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| **量化器** | 基础实现 | 弃疾 | 6月22日 | 6月25日 | 3 | P0 | +| | 预处理 | 莊霖 | | | 5 | P0 | +| | turbo分发 | 弃疾 | 6月26日 | 6月29日 | 2 | P0 | +| | FP32量化器 | 荐宁 | | | 3 | P0 | +| | FP16量化器 | 荐宁 | | | 3 | P0 | +| | Int8量化器 | 荐宁 | | | 3 | P0 | +| | Int4量化器 | 荐宁 | | | 3 | P0 | +| | PQ量化器 | 莊霖 | | | 10 | P1 | +| **索引对接** | HNSW | 荐宁 | | | 2 | P0 | +| | IVF | 荐宁 | | | 2 | P0 | +| | Vamana | 荐宁 | | | 2 | P0 | +| | DiskAnn | 弃疾 | 7月6日 | 7月8日 | 2 | P1 | +| **单元测试** | 单元测试 | ALL | | | 5 | P0 | +| **集成测试** | 集成测试 | ALL | | | 5 | P0 | + + +# 结语 ++ 本设计以**量化方案**为核心分类维度，将原本分散在量化器、度量、数据类型中的判断逻辑统一到 Turbo 量化器中。 ++ 确保一期聚焦高频使用场景的 Uniform/Record + INT8/INT4/FP16 + L2/Cosine/IP 组合，再逐步扩展 PQ量化器等复杂方案。 + + + +# 参考 ++ [距离计算能力](https://yuque.alibaba-inc.com/proxima/aby36h/lr69q2egk5ob18cw) ++ [CI&CD 支持集合](https://yuque.alibaba-inc.com/proxima/aby36h/pyvb301f3fs8s3ua) ++ [TurboQuant-近似最优失真率的在线量化方案](https://yuque.alibaba-inc.com/proxima/aby36h/lgllif6v3hht61be) diff --git a/doc/pq.md b/doc/pq.md new file mode 100644 index 000000000..58612e75d --- /dev/null +++ b/doc/pq.md @@ -0,0 +1,94 @@ +## 在turbo中实现preprocessor + +### 1. Preprocessor + +```cpp +class Preprocessor { + public: + using Pointer = std::shared_ptr; + + virtual ~Preprocessor() = default; + + //! Input dimensionality accepted by apply(). + virtual int in_dim() const = 0; + + //! Output dimensionality produced by apply(). May differ from in_dim() + //! (e.g. FHT pads to the next power of two). + virtual int out_dim() const = 0; + + //! Forward transform: map an input vector to the preprocessed space. + //! \p out must hold at least out_dim() elements. + virtual void apply(const float *in, float *out) const = 0; + + //! Inverse transform: recover the original-space vector from a preprocessed + //! one. \p out must hold at least in_dim() elements. + virtual void apply_inverse(const float *in, float *out) const = 0; + + //! Fit the preprocessor from a contiguous batch of training data. + //! \p data pointer to the first element of the batch. + //! \p num number of vectors in the batch. + //! \p stride byte offset between consecutive vectors (0 ⇒ packed). + virtual void train(const void *data, size_t num, size_t stride) = 0; + + //! Serialize the preprocessor into a self-contained blob (header + payload). + virtual int serialize(std::string *out) const = 0; + + //! Deserialize the preprocessor from a raw, possibly mmap-backed buffer. + virtual int deserialize(const void *data, size_t len) = 0; +}; +``` +定位：属于/src/turbo/quantizer的一部分，定义在/src/turbo/quantizer/preprocessor.h + +### 2. FhtRotator + +```cpp +class FhtRotator : public Preprocessor { + public: + using Pointer = std::shared_ptr; + + //! Kind of rotator. + virtual RotatorType type() const = 0; + + // in_dim(), out_dim(), apply(), apply_inverse(), train(), serialize(), + // and deserialize() are inherited from Preprocessor. +} + +//! Create an untrained rotator of the given type for in_dim dimensions. +FhtRotator::Pointer CreateFhtRotator(int in_dim); + +//! Create and restore a rotator from a serialized blob (reads the type from +//! the embedded RotatorSerHeader). Returns nullptr on malformed input. +FhtRotator::Pointer CreateFhtRotatorFromBlob(const void *data, size_t len); +``` + +### 3. 数据保存 + +#### 保存格式 +```cpp +struct RotatorSerHeader { + uint32_t magic; // kRotatorMagic + uint16_t version; // kRotatorSerVersion + uint16_t rotator_type; // RotatorType + uint32_t in_dim; // input dimensionality + uint32_t out_dim; // output dimensionality + uint32_t payload_size; // bytes following the header + uint32_t reserved; // 0, for future use / alignment +}; +static_assert(sizeof(RotatorSerHeader)) == 24, +``` +定义写在preprocessor.h中 + +在FHT中， +1. rotator_type设置为1，表示FHT旋转 +2. in_dim和out_dim设置为相同， +4. 具体实现参考/root/code/zvec/中fht的实现 + +#### 保存位置 +RotatorSerHeader放到当前已经存在的QuantizerSerHeader中，和他一起保存。具体来说： +1. 量化器调用serialize -> 其中调用preprocessor的serialize +2. serialize后的数据保存到Indexmeta中 +3. 当前方案已经有了保存Indexmeta，这样就可以自然保存 + + +### 4. 分发函数 +在 turbo.h/.cc 中新增函数，将fht的分发功能，也放在其中 diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 2eff572d5..c3506f68f 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -13,6 +13,8 @@ // limitations under the License. #pragma once +#include +#include #include #include @@ -33,6 +35,21 @@ using QueryPreprocessFunc = using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale, float bias, int8_t *out); +// FHT primitive function pointer types. +using FhtFlipSignFunc = void (*)(const uint8_t *flip, float *data, size_t dim); +using FhtKacsWalkFunc = void (*)(float *data, size_t len); +using FhtInplaceFunc = void (*)(float *data, size_t n); +using FhtVecRescaleFunc = void (*)(float *data, size_t n, float factor); + +// Aggregate of all FHT kernels needed by FhtRotator, dispatched by ISA. +struct FhtKernels { + FhtFlipSignFunc flip_sign; + FhtKacsWalkFunc kacs_walk; + FhtKacsWalkFunc inv_kacs_walk; + FhtInplaceFunc inplace; + FhtVecRescaleFunc rescale; +}; + enum class MetricType { kSquaredEuclidean, kCosine, @@ -90,4 +107,7 @@ QueryPreprocessFunc get_query_preprocess_func( // interface can grow to cover other output types (e.g. fp16) in the future. UniformQuantizeFunc get_uniform_quantize_func(DataType data_type); +// Returns all FHT kernels dispatched for the current CPU. +FhtKernels get_fht_kernels(); + } // namespace zvec::turbo diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index 7ec9c0e5d..f34f24e66 100644 --- a/src/turbo/CMakeLists.txt +++ b/src/turbo/CMakeLists.txt @@ -13,15 +13,32 @@ endif() file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h) -# Set per-file compile flags for AVX512-VNNI sources. +# Set per-file compile flags for SIMD sources. # set_source_files_properties is directory-scoped, so it must be called in the # same directory that adds the sources to a target (i.e. here, not in a # subdirectory). if(NOT ANDROID AND AUTO_DETECT_ARCH) if (HOST_ARCH MATCHES "^(x86|x64)$") - file(GLOB_RECURSE AVX512_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc) + # SSE + file(GLOB_RECURSE SSE_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/sse/*.cc) set_source_files_properties( - ${AVX512_VNNI_SRCS} + ${SSE_SRCS} + PROPERTIES + COMPILE_FLAGS "${TURBO_MARCH_FLAG_SSE}" + ) + + # AVX2 + file(GLOB_RECURSE AVX2_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx2/*.cc) + set_source_files_properties( + ${AVX2_SRCS} + PROPERTIES + COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX2}" + ) + + # AVX512 (avx512_vnni and avx512/fht both use AVX512 march) + file(GLOB_RECURSE AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512/*.cc) + set_source_files_properties( + ${AVX512_SRCS} PROPERTIES COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}" ) diff --git a/src/turbo/distance/avx2/fht/fht.h b/src/turbo/distance/avx2/fht/fht.h new file mode 100644 index 000000000..5e636a54f --- /dev/null +++ b/src/turbo/distance/avx2/fht/fht.h @@ -0,0 +1,37 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::avx2 { + +//! Apply bitwise sign-flip mask to a float vector (AVX2). +void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation (AVX2). +void fht_kacs_walk_avx2(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation (AVX2). +void fht_inv_kacs_walk_avx2(float *data, size_t len); + +//! In-place Fast Hadamard Transform (AVX2, n must be power-of-2). +void fht_inplace_avx2(float *data, size_t n); + +//! Element-wise rescale: data[i] *= factor (AVX2). +void fht_vec_rescale_avx2(float *data, size_t n, float factor); + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/fht/fht_avx2.cc b/src/turbo/distance/avx2/fht/fht_avx2.cc new file mode 100644 index 000000000..903305931 --- /dev/null +++ b/src/turbo/distance/avx2/fht/fht_avx2.cc @@ -0,0 +1,138 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__AVX2__) + +#include "fht.h" + +#include +#include +#include +#include +#include + +namespace zvec::turbo::avx2 { + +void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~31u; + constexpr size_t kChunk = 32; + const __m256i bit_select = + _mm256_setr_epi32(0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80); + const __m256 sign_flip = _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000)); + for (size_t i = 0; i < simd_end; i += kChunk) { + uint32_t mask_bits; + std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits)); + for (int b = 0; b < 4; ++b) { + __m256i mb = _mm256_set1_epi32((mask_bits >> (b * 8)) & 0xFF); + __m256i test = _mm256_and_si256(mb, bit_select); + __m256i cmp = _mm256_cmpeq_epi32(test, bit_select); + __m256 xor_mask = _mm256_and_ps(_mm256_castsi256_ps(cmp), sign_flip); + __m256 v = _mm256_loadu_ps(&data[i + b * 8]); + v = _mm256_xor_ps(v, xor_mask); + _mm256_storeu_ps(&data[i + b * 8], v); + } + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_avx2(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~7u; + for (size_t i = 0; i < half_end; i += 8) { + __m256 x = _mm256_loadu_ps(&data[i]); + __m256 y = _mm256_loadu_ps(&data[i + offset]); + _mm256_storeu_ps(&data[i], _mm256_add_ps(x, y)); + _mm256_storeu_ps(&data[i + offset], _mm256_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_avx2(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~7u; + const __m256 half_fac = _mm256_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 8) { + __m256 a = _mm256_loadu_ps(&data[i]); + __m256 b = _mm256_loadu_ps(&data[i + offset]); + _mm256_storeu_ps(&data[i], _mm256_mul_ps(_mm256_add_ps(a, b), half_fac)); + _mm256_storeu_ps(&data[i + offset], + _mm256_mul_ps(_mm256_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace_avx2(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + size_t step = len << 1; + size_t simd_end = len & ~7u; + for (size_t i = 0; i < n; i += step) { + for (size_t j = 0; j < simd_end; j += 8) { + __m256 u = _mm256_loadu_ps(&data[i + j]); + __m256 v = _mm256_loadu_ps(&data[i + j + len]); + _mm256_storeu_ps(&data[i + j], _mm256_add_ps(u, v)); + _mm256_storeu_ps(&data[i + j + len], _mm256_sub_ps(u, v)); + } + for (size_t j = simd_end; j < len; ++j) { + float u = data[i + j]; + float v = data[i + j + len]; + data[i + j] = u + v; + data[i + j + len] = u - v; + } + } + } +} + +void fht_vec_rescale_avx2(float *data, size_t n, float factor) { + const __m256 fac = _mm256_set1_ps(factor); + size_t simd_end = n & ~7u; + for (size_t i = 0; i < simd_end; i += 8) { + __m256 v = _mm256_loadu_ps(&data[i]); + _mm256_storeu_ps(&data[i], _mm256_mul_ps(v, fac)); + } + // Scalar tail + for (size_t i = simd_end; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::avx2 + +#endif // __AVX2__ diff --git a/src/turbo/distance/avx512/fht/fht.h b/src/turbo/distance/avx512/fht/fht.h new file mode 100644 index 000000000..53241a0c5 --- /dev/null +++ b/src/turbo/distance/avx512/fht/fht.h @@ -0,0 +1,37 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::avx512 { + +//! Apply bitwise sign-flip mask to a float vector (AVX512). +void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation (AVX512). +void fht_kacs_walk_avx512(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation (AVX512). +void fht_inv_kacs_walk_avx512(float *data, size_t len); + +//! In-place Fast Hadamard Transform (AVX512, n must be power-of-2). +void fht_inplace_avx512(float *data, size_t n); + +//! Element-wise rescale: data[i] *= factor (AVX512). +void fht_vec_rescale_avx512(float *data, size_t n, float factor); + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/fht/fht_avx512.cc b/src/turbo/distance/avx512/fht/fht_avx512.cc new file mode 100644 index 000000000..b7ea37d63 --- /dev/null +++ b/src/turbo/distance/avx512/fht/fht_avx512.cc @@ -0,0 +1,143 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__AVX512F__) + +#include "fht.h" + +#include +#include +#include +#include +#include + +namespace zvec::turbo::avx512 { + +void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~63u; + constexpr size_t kChunk = 64; + const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000)); + for (size_t i = 0; i < simd_end; i += kChunk) { + uint64_t mask_bits; + std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits)); + const __mmask16 m0 = _cvtu32_mask16(mask_bits & 0xFFFF); + const __mmask16 m1 = _cvtu32_mask16((mask_bits >> 16) & 0xFFFF); + const __mmask16 m2 = _cvtu32_mask16((mask_bits >> 32) & 0xFFFF); + const __mmask16 m3 = _cvtu32_mask16((mask_bits >> 48) & 0xFFFF); + __m512 v0 = _mm512_loadu_ps(&data[i]); + v0 = _mm512_mask_xor_ps(v0, m0, v0, sign_flip); + _mm512_storeu_ps(&data[i], v0); + __m512 v1 = _mm512_loadu_ps(&data[i + 16]); + v1 = _mm512_mask_xor_ps(v1, m1, v1, sign_flip); + _mm512_storeu_ps(&data[i + 16], v1); + __m512 v2 = _mm512_loadu_ps(&data[i + 32]); + v2 = _mm512_mask_xor_ps(v2, m2, v2, sign_flip); + _mm512_storeu_ps(&data[i + 32], v2); + __m512 v3 = _mm512_loadu_ps(&data[i + 48]); + v3 = _mm512_mask_xor_ps(v3, m3, v3, sign_flip); + _mm512_storeu_ps(&data[i + 48], v3); + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_avx512(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~15u; + for (size_t i = 0; i < half_end; i += 16) { + __m512 x = _mm512_loadu_ps(&data[i]); + __m512 y = _mm512_loadu_ps(&data[i + offset]); + _mm512_storeu_ps(&data[i], _mm512_add_ps(x, y)); + _mm512_storeu_ps(&data[i + offset], _mm512_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_avx512(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~15u; + const __m512 half_fac = _mm512_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 16) { + __m512 a = _mm512_loadu_ps(&data[i]); + __m512 b = _mm512_loadu_ps(&data[i + offset]); + _mm512_storeu_ps(&data[i], _mm512_mul_ps(_mm512_add_ps(a, b), half_fac)); + _mm512_storeu_ps(&data[i + offset], + _mm512_mul_ps(_mm512_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace_avx512(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + size_t step = len << 1; + size_t simd_end = len & ~15u; + for (size_t i = 0; i < n; i += step) { + for (size_t j = 0; j < simd_end; j += 16) { + __m512 u = _mm512_loadu_ps(&data[i + j]); + __m512 v = _mm512_loadu_ps(&data[i + j + len]); + _mm512_storeu_ps(&data[i + j], _mm512_add_ps(u, v)); + _mm512_storeu_ps(&data[i + j + len], _mm512_sub_ps(u, v)); + } + for (size_t j = simd_end; j < len; ++j) { + float u = data[i + j]; + float v = data[i + j + len]; + data[i + j] = u + v; + data[i + j + len] = u - v; + } + } + } +} + +void fht_vec_rescale_avx512(float *data, size_t n, float factor) { + const __m512 fac = _mm512_set1_ps(factor); + size_t simd_end = n & ~15u; + for (size_t i = 0; i < simd_end; i += 16) { + __m512 v = _mm512_loadu_ps(&data[i]); + _mm512_storeu_ps(&data[i], _mm512_mul_ps(v, fac)); + } + // Scalar tail + for (size_t i = simd_end; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::avx512 + +#endif // __AVX512F__ diff --git a/src/turbo/distance/scalar/fht/fht.h b/src/turbo/distance/scalar/fht/fht.h new file mode 100644 index 000000000..7ba98be5d --- /dev/null +++ b/src/turbo/distance/scalar/fht/fht.h @@ -0,0 +1,38 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::scalar { + +//! Apply bitwise sign-flip mask to a float vector. +//! Each bit in \p flip controls one element: bit set means negate. +void fht_flip_sign(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation to non-power-of-2 FHT. +void fht_kacs_walk(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation. +void fht_inv_kacs_walk(float *data, size_t len); + +//! In-place Fast Hadamard Transform on \p n elements (must be power-of-2). +void fht_inplace(float *data, size_t n); + +//! Element-wise rescale: data[i] *= factor. +void fht_vec_rescale(float *data, size_t n, float factor); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fht/fht_scalar.cc b/src/turbo/distance/scalar/fht/fht_scalar.cc new file mode 100644 index 000000000..008c6a2a6 --- /dev/null +++ b/src/turbo/distance/scalar/fht/fht_scalar.cc @@ -0,0 +1,80 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "fht.h" + +#include +#include +#include + +namespace zvec::turbo::scalar { + +void fht_flip_sign(const uint8_t *flip, float *data, size_t dim) { + for (size_t i = 0; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + for (size_t i = 0; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + for (size_t i = 0; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_inplace(float *data, size_t n) { + for (size_t len = 1; len < n; len <<= 1) { + for (size_t i = 0; i < n; i += len << 1) { + for (size_t j = i; j < i + len; ++j) { + float u = data[j]; + float v = data[j + len]; + data[j] = u + v; + data[j + len] = u - v; + } + } + } +} + +void fht_vec_rescale(float *data, size_t n, float factor) { + for (size_t i = 0; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/sse/fht/fht.h b/src/turbo/distance/sse/fht/fht.h new file mode 100644 index 000000000..9e1cd20a5 --- /dev/null +++ b/src/turbo/distance/sse/fht/fht.h @@ -0,0 +1,34 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::sse { + +//! Apply bitwise sign-flip mask to a float vector (SSE). +void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim); + +//! Apply KacsWalk butterfly operation (SSE). +void fht_kacs_walk_sse(float *data, size_t len); + +//! Inverse KacsWalk butterfly operation (SSE). +void fht_inv_kacs_walk_sse(float *data, size_t len); + +//! Element-wise rescale: data[i] *= factor (SSE). +void fht_vec_rescale_sse(float *data, size_t n, float factor); + +} // namespace zvec::turbo::sse diff --git a/src/turbo/distance/sse/fht/fht_sse.cc b/src/turbo/distance/sse/fht/fht_sse.cc new file mode 100644 index 000000000..7dbbf1951 --- /dev/null +++ b/src/turbo/distance/sse/fht/fht_sse.cc @@ -0,0 +1,119 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__SSE2__) + +#include "fht.h" + +#include +#include +#include +#include +#include + +namespace zvec::turbo::sse { + +void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~3u; + size_t flip_bytes = (dim + 7) / 8; + for (size_t i = 0; i < simd_end; i += 4) { + uint16_t bits16; + size_t byte_pos = i / 8; + if (byte_pos + 1 < flip_bytes) { + std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16)); + } else { + bits16 = flip[byte_pos]; + } + bits16 >>= (i % 8); + uint32_t b0 = bits16 & 1u; + uint32_t b1 = (bits16 >> 1) & 1u; + uint32_t b2 = (bits16 >> 2) & 1u; + uint32_t b3 = (bits16 >> 3) & 1u; + __m128i bit_mask = _mm_set_epi32(b3, b2, b1, b0); + __m128i sign_mask = _mm_slli_epi32(bit_mask, 31); + __m128 v = _mm_loadu_ps(&data[i]); + v = _mm_xor_ps(v, _mm_castsi128_ps(sign_mask)); + _mm_storeu_ps(&data[i], v); + } + // Scalar tail + for (size_t i = simd_end; i < dim; ++i) { + if (flip[i / 8] & (1u << (i % 8))) { + data[i] = -data[i]; + } + } +} + +void fht_kacs_walk_sse(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + size_t half_end = half & ~3u; + for (size_t i = 0; i < half_end; i += 4) { + __m128 x = _mm_loadu_ps(&data[i]); + __m128 y = _mm_loadu_ps(&data[i + offset]); + _mm_storeu_ps(&data[i], _mm_add_ps(x, y)); + _mm_storeu_ps(&data[i + offset], _mm_sub_ps(x, y)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float x = data[i]; + float y = data[i + offset]; + data[i] = x + y; + data[i + offset] = x - y; + } + if (base != 0) { + data[half] *= std::sqrt(2.0f); + } +} + +void fht_inv_kacs_walk_sse(float *data, size_t len) { + size_t half = len / 2; + size_t base = len % 2; + size_t offset = base + half; + if (base != 0) { + data[half] *= std::sqrt(0.5f); + } + size_t half_end = half & ~3u; + const __m128 half_fac = _mm_set1_ps(0.5f); + for (size_t i = 0; i < half_end; i += 4) { + __m128 a = _mm_loadu_ps(&data[i]); + __m128 b = _mm_loadu_ps(&data[i + offset]); + _mm_storeu_ps(&data[i], _mm_mul_ps(_mm_add_ps(a, b), half_fac)); + _mm_storeu_ps(&data[i + offset], _mm_mul_ps(_mm_sub_ps(a, b), half_fac)); + } + // Scalar tail + for (size_t i = half_end; i < half; ++i) { + float a = data[i]; + float b = data[i + offset]; + data[i] = (a + b) * 0.5f; + data[i + offset] = (a - b) * 0.5f; + } +} + +void fht_vec_rescale_sse(float *data, size_t n, float factor) { + const __m128 fac = _mm_set1_ps(factor); + size_t simd_end = n & ~3u; + for (size_t i = 0; i < simd_end; i += 4) { + __m128 v = _mm_loadu_ps(&data[i]); + _mm_storeu_ps(&data[i], _mm_mul_ps(v, fac)); + } + // Scalar tail + for (size_t i = simd_end; i < n; ++i) { + data[i] *= factor; + } +} + +} // namespace zvec::turbo::sse + +#endif // __SSE2__ diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.cc b/src/turbo/quantizer/preprocessor/fht_rotator.cc new file mode 100644 index 000000000..94a19a916 --- /dev/null +++ b/src/turbo/quantizer/preprocessor/fht_rotator.cc @@ -0,0 +1,266 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "fht_rotator.h" + +#include +#include +#include +#include + +namespace zvec { +namespace turbo { + +using core::IndexError_InvalidArgument; +using core::IndexError_Runtime; +using core::IndexError_Unsupported; + +// ============================================================================ +// FhtRotator method implementations +// ============================================================================ + +size_t FhtRotator::floor_pow2(size_t n) { + if (n == 0) return 0; + size_t p = 1; + while (p * 2 <= n) p *= 2; + return p; +} + +FhtRotator::Pointer FhtRotator::create(int dim) { + if (dim <= 0) return nullptr; + + Pointer r(new FhtRotator()); + r->in_dim_ = dim; + r->out_dim_ = dim; + r->trunc_dim_ = floor_pow2(static_cast(dim)); + r->fac_ = 1.0f / std::sqrt(static_cast(r->trunc_dim_)); + r->flip_offset_ = (static_cast(dim) + kByteLen - 1) / kByteLen; + auto k = get_fht_kernels(); + r->flip_sign_fn_ = k.flip_sign; + r->kacs_walk_fn_ = k.kacs_walk; + r->inv_kacs_walk_fn_ = k.inv_kacs_walk; + r->inplace_fn_ = k.inplace; + r->rescale_fn_ = k.rescale; + // flip_ stays empty until train() is called. + return r; +} + +FhtRotator::Pointer FhtRotator::from_blob(const void *data, size_t len) { + if (!data || len < sizeof(RotatorSerHeader)) return nullptr; + + const auto *hdr = reinterpret_cast(data); + if (hdr->magic != kRotatorMagic) return nullptr; + if (hdr->version != kRotatorSerVersion) return nullptr; + if (static_cast(hdr->rotator_type) != RotatorType::kFht) { + return nullptr; + } + + Pointer r(new FhtRotator()); + const size_t expected_total = + sizeof(RotatorSerHeader) + static_cast(hdr->payload_size); + if (len < expected_total) return nullptr; + + int rc = r->deserialize(data, len); + if (rc != 0) return nullptr; + return r; +} + +void FhtRotator::train(const void * /*data*/, size_t /*num*/, + size_t /*stride*/) { + if (in_dim_ <= 0) return; + + flip_offset_ = (static_cast(in_dim_) + kByteLen - 1) / kByteLen; + flip_.resize(4 * flip_offset_); + + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_int_distribution dist(0, 255); + for (auto &b : flip_) b = static_cast(dist(gen)); +} + +// --------------------------------------------------------------------------- +// apply (forward rotation) +// --------------------------------------------------------------------------- + +void FhtRotator::apply(const float *in, float *out) const { + const size_t dim = static_cast(in_dim_); + std::memcpy(out, in, sizeof(float) * dim); + + if (trunc_dim_ == dim) { + // Exact power-of-2: 4 rounds of (flip -> FHT -> rescale) + flip_sign_fn_(flip_.data(), out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + flip_sign_fn_(flip_.data() + flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + flip_sign_fn_(flip_.data() + 2 * flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + flip_sign_fn_(flip_.data() + 3 * flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + + return; + } + + // Non-power-of-2: 4 rounds with kacs_walk + size_t start = dim - trunc_dim_; + float *trunc_ptr = out + start; + + // Round 1: FHT on [0, trunc_dim) + flip_sign_fn_(flip_.data(), out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Round 2: FHT on [start, start + trunc_dim) + flip_sign_fn_(flip_.data() + flip_offset_, out, dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Round 3: FHT on [0, trunc_dim) + flip_sign_fn_(flip_.data() + 2 * flip_offset_, out, dim); + inplace_fn_(out, trunc_dim_); + rescale_fn_(out, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Round 4: FHT on [start, start + trunc_dim) + flip_sign_fn_(flip_.data() + 3 * flip_offset_, out, dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, fac_); + kacs_walk_fn_(out, dim); + + // Final rescale: combine the 4 kacs_walk reductions + rescale_fn_(out, dim, 0.25f); +} + +// --------------------------------------------------------------------------- +// apply_inverse (inverse rotation) +// --------------------------------------------------------------------------- + +void FhtRotator::apply_inverse(const float *in, float *out) const { + const size_t dim = static_cast(in_dim_); + // Copy input into working buffer + std::vector data(in, in + dim); + + if (trunc_dim_ == dim) { + // Exact power-of-2: reverse 4 rounds in reverse order. + const float inv_fac = 1.0f / std::sqrt(static_cast(trunc_dim_)); + for (int round = 3; round >= 0; --round) { + inplace_fn_(data.data(), trunc_dim_); + rescale_fn_(data.data(), trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + static_cast(round) * flip_offset_, + data.data(), dim); + } + std::memcpy(out, data.data(), dim * sizeof(float)); + return; + } + + // Non-power-of-2: undo final rescale(0.25) first + rescale_fn_(data.data(), dim, 4.0f); + + const float inv_fac = 1.0f / std::sqrt(static_cast(trunc_dim_)); + size_t start = dim - trunc_dim_; + float *trunc_ptr = data.data() + start; + + // Undo Round 4 (FHT on [start, start+trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + 3 * flip_offset_, data.data(), dim); + + // Undo Round 3 (FHT on [0, trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(data.data(), trunc_dim_); + rescale_fn_(data.data(), trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + 2 * flip_offset_, data.data(), dim); + + // Undo Round 2 (FHT on [start, start+trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(trunc_ptr, trunc_dim_); + rescale_fn_(trunc_ptr, trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data() + flip_offset_, data.data(), dim); + + // Undo Round 1 (FHT on [0, trunc_dim)) + inv_kacs_walk_fn_(data.data(), dim); + inplace_fn_(data.data(), trunc_dim_); + rescale_fn_(data.data(), trunc_dim_, inv_fac); + flip_sign_fn_(flip_.data(), data.data(), dim); + + std::memcpy(out, data.data(), dim * sizeof(float)); +} + +// --------------------------------------------------------------------------- +// serialize / deserialize +// --------------------------------------------------------------------------- + +int FhtRotator::serialize(std::string *out) const { + if (!out) return IndexError_InvalidArgument; + if (flip_.empty()) return IndexError_Runtime; + + RotatorSerHeader hdr{}; + hdr.magic = kRotatorMagic; + hdr.version = kRotatorSerVersion; + hdr.rotator_type = static_cast(RotatorType::kFht); + hdr.in_dim = static_cast(in_dim_); + hdr.out_dim = static_cast(out_dim_); + hdr.payload_size = static_cast(flip_.size()); + hdr.reserved = 0; + + out->resize(sizeof(hdr) + flip_.size()); + std::memcpy(&(*out)[0], &hdr, sizeof(hdr)); + std::memcpy(&(*out)[sizeof(hdr)], flip_.data(), flip_.size()); + return 0; +} + +int FhtRotator::deserialize(const void *data, size_t len) { + if (!data || len < sizeof(RotatorSerHeader)) return IndexError_InvalidArgument; + + const auto *hdr = reinterpret_cast(data); + if (hdr->magic != kRotatorMagic) return IndexError_Unsupported; + if (hdr->version != kRotatorSerVersion) return IndexError_Unsupported; + if (static_cast(hdr->rotator_type) != RotatorType::kFht) { + return IndexError_Unsupported; + } + + const size_t total = sizeof(RotatorSerHeader) + hdr->payload_size; + if (len < total) return IndexError_InvalidArgument; + + in_dim_ = static_cast(hdr->in_dim); + out_dim_ = static_cast(hdr->out_dim); + trunc_dim_ = floor_pow2(static_cast(in_dim_)); + fac_ = 1.0f / std::sqrt(static_cast(trunc_dim_)); + flip_offset_ = (static_cast(in_dim_) + kByteLen - 1) / kByteLen; + auto k = get_fht_kernels(); + flip_sign_fn_ = k.flip_sign; + kacs_walk_fn_ = k.kacs_walk; + inv_kacs_walk_fn_ = k.inv_kacs_walk; + inplace_fn_ = k.inplace; + rescale_fn_ = k.rescale; + + flip_.resize(hdr->payload_size); + std::memcpy(flip_.data(), + reinterpret_cast(data) + sizeof(RotatorSerHeader), + hdr->payload_size); + + return 0; +} + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/preprocessor/fht_rotator.h b/src/turbo/quantizer/preprocessor/fht_rotator.h new file mode 100644 index 000000000..15dfbb0c4 --- /dev/null +++ b/src/turbo/quantizer/preprocessor/fht_rotator.h @@ -0,0 +1,97 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include "preprocessor.h" +#include + +namespace zvec { +namespace turbo { + +// ============================================================================ +// FhtRotator - O(d log d) FHT-based Kac random rotation +// +// Works with any dimension (non-power-of-2 uses trunc_dim + KacsWalk). +// When dimension is a power of 2, uses 4 rounds of (flip -> FHT -> rescale). +// When dimension is NOT a power of 2, uses kacs_walk reduction. +// ============================================================================ + +class FhtRotator : public Preprocessor { + public: + using Pointer = std::shared_ptr; + + //! Create an untrained rotator for \p in_dim dimensions. + //! Call train() afterwards to generate the random flip-sign arrays. + static Pointer create(int in_dim); + + //! Create and restore a rotator from a serialized blob (reads the type from + //! the embedded RotatorSerHeader). Returns nullptr on malformed input. + static Pointer from_blob(const void *data, size_t len); + + // -- Preprocessor interface ------------------------------------------------ + + int in_dim() const override { return in_dim_; } + int out_dim() const override { return out_dim_; } + + void apply(const float *in, float *out) const override; + void apply_inverse(const float *in, float *out) const override; + + //! Generate 4 rounds of random flip-sign arrays. + //! For FhtRotator the training data is ignored; only \p in_dim() matters. + void train(const void *data, size_t num, size_t stride) override; + + int serialize(std::string *out) const override; + int deserialize(const void *data, size_t len) override; + + //! Rotator type tag (kFht = 1). + RotatorType rotator_type() const { return RotatorType::kFht; } + + private: + FhtRotator() = default; + + //! Largest power of 2 <= dim. + static size_t floor_pow2(size_t n); + + int in_dim_{0}; + int out_dim_{0}; + + //! Packed flip-sign bits: 4 rounds, each ceil(in_dim / 8) bytes. + std::vector flip_; + + //! Bytes per round: ceil(in_dim / 8). + size_t flip_offset_{0}; + + //! Largest power of 2 <= in_dim (used for FHT length). + size_t trunc_dim_{0}; + + //! Rescale factor: 1 / sqrt(trunc_dim). + float fac_{0}; + + //! ISA-dispatched FHT kernels (cached at construction time). + FhtFlipSignFunc flip_sign_fn_{nullptr}; + FhtKacsWalkFunc kacs_walk_fn_{nullptr}; + FhtKacsWalkFunc inv_kacs_walk_fn_{nullptr}; + FhtInplaceFunc inplace_fn_{nullptr}; + FhtVecRescaleFunc rescale_fn_{nullptr}; + + static constexpr size_t kByteLen = 8; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/preprocessor/preprocessor.h b/src/turbo/quantizer/preprocessor/preprocessor.h new file mode 100644 index 000000000..07ed313ea --- /dev/null +++ b/src/turbo/quantizer/preprocessor/preprocessor.h @@ -0,0 +1,93 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +namespace zvec { +namespace turbo { + +//! Magic number ('ROTR') stamped at the start of a serialized rotator blob. +constexpr uint32_t kRotatorMagic = 0x52544F52u; +//! Current rotator serialization format version. +constexpr uint16_t kRotatorSerVersion = 1; + +//! Kind of rotator (stored in RotatorSerHeader::rotator_type). +enum class RotatorType : uint16_t { + kFht = 0, //!< O(d log d) FHT-based Kac random rotation +}; + +//! Self-describing, fixed-size header that prefixes every serialized rotator. +//! The type-specific payload (flip signs, rotation matrix, ...) follows +//! immediately after this header. +struct RotatorSerHeader { + uint32_t magic; // kRotatorMagic + uint16_t version; // kRotatorSerVersion + uint16_t rotator_type; // RotatorType + uint32_t in_dim; // input dimensionality + uint32_t out_dim; // output dimensionality + uint32_t payload_size; // bytes following the header + uint32_t reserved; // 0, for future use / alignment +}; +static_assert(sizeof(RotatorSerHeader) == 24, + "RotatorSerHeader must be 24 bytes"); + +//! Abstract preprocessor interface. +//! +//! A Preprocessor applies a deterministic, invertible transform to each +//! vector (e.g. random rotation). Concrete subclasses (FhtRotator, ...) +//! implement the actual algorithm. +class Preprocessor { + public: + using Pointer = std::shared_ptr; + + virtual ~Preprocessor() = default; + + //! Input dimensionality accepted by apply(). + virtual int in_dim() const = 0; + + //! Output dimensionality produced by apply(). May differ from in_dim() + //! (e.g. FHT pads to the next power of two), but for FhtRotator they are + //! equal. + virtual int out_dim() const = 0; + + //! Forward transform: map an input vector to the preprocessed space. + //! \p out must hold at least out_dim() elements. + virtual void apply(const float *in, float *out) const = 0; + + //! Inverse transform: recover the original-space vector from a preprocessed + //! one. \p out must hold at least in_dim() elements. + virtual void apply_inverse(const float *in, float *out) const = 0; + + //! Fit / initialize the preprocessor from a contiguous batch of training + //! data. For FhtRotator this generates the random flip-sign arrays. + //! \p data pointer to the first element of the batch. + //! \p num number of vectors in the batch. + //! \p stride byte offset between consecutive vectors (0 => packed). + virtual void train(const void *data, size_t num, size_t stride) = 0; + + //! Serialize the preprocessor into a self-contained blob + //! (RotatorSerHeader + payload). + virtual int serialize(std::string *out) const = 0; + + //! Deserialize the preprocessor from a raw, possibly mmap-backed buffer. + virtual int deserialize(const void *data, size_t len) = 0; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 80f2ee77d..8e2e1b617 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -18,10 +18,21 @@ #include "avx512_vnni/record_quantized_int8/squared_euclidean.h" #include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/squared_euclidean.h" +#include "scalar/fht/fht.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" +#if defined(__SSE2__) +#include "sse/fht/fht.h" +#endif +#if defined(__AVX2__) +#include "avx2/fht/fht.h" +#endif +#if defined(__AVX512F__) +#include "avx512/fht/fht.h" +#endif + namespace zvec::turbo { DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, @@ -143,4 +154,47 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } +FhtKernels get_fht_kernels() { + FhtKernels k; + // Default: scalar fallback for all + k.flip_sign = scalar::fht_flip_sign; + k.kacs_walk = scalar::fht_kacs_walk; + k.inv_kacs_walk = scalar::fht_inv_kacs_walk; + k.inplace = scalar::fht_inplace; + k.rescale = scalar::fht_vec_rescale; + +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ) { + k.flip_sign = avx512::fht_flip_sign_avx512; + k.kacs_walk = avx512::fht_kacs_walk_avx512; + k.inv_kacs_walk = avx512::fht_inv_kacs_walk_avx512; + k.inplace = avx512::fht_inplace_avx512; + k.rescale = avx512::fht_vec_rescale_avx512; + return k; + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) { + k.flip_sign = avx2::fht_flip_sign_avx2; + k.kacs_walk = avx2::fht_kacs_walk_avx2; + k.inv_kacs_walk = avx2::fht_inv_kacs_walk_avx2; + k.inplace = avx2::fht_inplace_avx2; + k.rescale = avx2::fht_vec_rescale_avx2; + return k; + } +#endif +#if defined(__SSE2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2) { + k.flip_sign = sse::fht_flip_sign_sse; + k.kacs_walk = sse::fht_kacs_walk_sse; + k.inv_kacs_walk = sse::fht_inv_kacs_walk_sse; + k.rescale = sse::fht_vec_rescale_sse; + // inplace fallback to scalar (SSE has no fht_inplace) + return k; + } +#endif + return k; // scalar +} + } // namespace zvec::turbo diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc new file mode 100644 index 000000000..178dad3fd --- /dev/null +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -0,0 +1,204 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include "quantizer/preprocessor/fht_rotator.h" + +using namespace zvec::turbo; + +namespace { + +// Helper: fill a vector with random floats. +void fill_random(float *data, size_t dim, std::mt19937 &gen) { + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (size_t i = 0; i < dim; ++i) data[i] = dist(gen); +} + +// Helper: check round-trip (apply_inverse(apply(x)) == x) within tolerance. +void check_round_trip(const FhtRotator &rot, const std::vector &input, + float tol = 1e-3f) { + const int dim = rot.in_dim(); + std::vector rotated(dim); + std::vector recovered(dim); + + rot.apply(input.data(), rotated.data()); + rot.apply_inverse(rotated.data(), recovered.data()); + + for (int i = 0; i < dim; ++i) { + EXPECT_NEAR(input[i], recovered[i], tol) + << "mismatch at i=" << i << " dim=" << dim; + } +} + +} // anonymous namespace + +// --------------------------------------------------------------------------- +// Power-of-2 dimensions +// --------------------------------------------------------------------------- + +TEST(FhtRotator, PowerOf2RoundTrip) { + std::mt19937 gen(42); + for (int dim : {1, 2, 4, 8, 16, 32, 64, 128, 256}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot) << "create failed for dim=" << dim; + rot->train(nullptr, 0, 0); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + check_round_trip(*rot, input); + } +} + +// --------------------------------------------------------------------------- +// Non-power-of-2 dimensions +// --------------------------------------------------------------------------- + +TEST(FhtRotator, NonPowerOf2RoundTrip) { + std::mt19937 gen(123); + for (int dim : {3, 5, 7, 10, 13, 31, 50, 97, 100, 127, 192, 320}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot) << "create failed for dim=" << dim; + rot->train(nullptr, 0, 0); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + check_round_trip(*rot, input); + } +} + +// --------------------------------------------------------------------------- +// Serialize / Deserialize round-trip +// --------------------------------------------------------------------------- + +TEST(FhtRotator, SerializeDeserialize) { + std::mt19937 gen(999); + for (int dim : {32, 97, 128}) { + // Build and train original rotator. + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + // Serialize. + std::string blob; + ASSERT_EQ(0, rot->serialize(&blob)); + ASSERT_GT(blob.size(), sizeof(RotatorSerHeader)); + + // Restore from blob. + auto rot2 = FhtRotator::from_blob(blob.data(), blob.size()); + ASSERT_TRUE(rot2) << "from_blob failed for dim=" << dim; + + // Dimensions must match. + EXPECT_EQ(rot2->in_dim(), dim); + EXPECT_EQ(rot2->out_dim(), dim); + + // Round-trip via the restored rotator must produce the same result + // as the original (same flip signs). + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector r1(dim), r2(dim); + rot->apply(input.data(), r1.data()); + rot2->apply(input.data(), r2.data()); + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) << "apply mismatch at i=" << i; + } + + // Inverse via restored rotator must recover the input. + check_round_trip(*rot2, input); + } +} + +// --------------------------------------------------------------------------- +// Dimension preserved +// --------------------------------------------------------------------------- + +TEST(FhtRotator, DimensionPreserved) { + for (int dim : {1, 7, 64, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + EXPECT_EQ(rot->in_dim(), dim); + EXPECT_EQ(rot->out_dim(), dim); + } +} + +// --------------------------------------------------------------------------- +// Train generates non-zero flip signs +// --------------------------------------------------------------------------- + +TEST(FhtRotator, TrainGeneratesFlip) { + for (int dim : {8, 64, 97}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + // Before train, apply should not crash but flip is empty — we skip + // calling apply before train. After train, serialize must succeed + // (which requires flip to be populated). + rot->train(nullptr, 0, 0); + + std::string blob; + EXPECT_EQ(0, rot->serialize(&blob)) + << "serialize failed after train for dim=" << dim; + + // Verify the payload is not all zeros (extremely unlikely for random bits). + const auto *hdr = reinterpret_cast(blob.data()); + EXPECT_EQ(hdr->magic, kRotatorMagic); + EXPECT_EQ(hdr->version, kRotatorSerVersion); + EXPECT_EQ(static_cast(hdr->rotator_type), RotatorType::kFht); + EXPECT_EQ(static_cast(hdr->in_dim), dim); + EXPECT_EQ(static_cast(hdr->out_dim), dim); + EXPECT_GT(hdr->payload_size, 0u); + + // Check that the flip payload is not all-zero. + const uint8_t *payload = + reinterpret_cast(blob.data()) + sizeof(RotatorSerHeader); + bool any_nonzero = false; + for (uint32_t i = 0; i < hdr->payload_size; ++i) { + if (payload[i] != 0) { any_nonzero = true; break; } + } + EXPECT_TRUE(any_nonzero) << "flip payload is all-zero for dim=" << dim; + } +} + +// --------------------------------------------------------------------------- +// Create with invalid dimension returns nullptr +// --------------------------------------------------------------------------- + +TEST(FhtRotator, InvalidDimension) { + EXPECT_EQ(FhtRotator::create(0), nullptr); + EXPECT_EQ(FhtRotator::create(-1), nullptr); +} + +// --------------------------------------------------------------------------- +// from_blob with malformed input returns nullptr +// --------------------------------------------------------------------------- + +TEST(FhtRotator, FromBlobMalformed) { + EXPECT_EQ(FhtRotator::from_blob(nullptr, 0), nullptr); + + // Too short. + char buf[4] = {}; + EXPECT_EQ(FhtRotator::from_blob(buf, sizeof(buf)), nullptr); + + // Wrong magic. + RotatorSerHeader hdr{}; + hdr.magic = 0xDEADBEEF; + hdr.version = kRotatorSerVersion; + hdr.rotator_type = static_cast(RotatorType::kFht); + hdr.payload_size = 0; + EXPECT_EQ(FhtRotator::from_blob(&hdr, sizeof(hdr)), nullptr); +} From 437a7546c7df926ab748ce129f2cfa24083b0edc Mon Sep 17 00:00:00 2001 From: ray Date: Thu, 2 Jul 2026 16:39:13 +0800 Subject: [PATCH 18/45] fix: fix symbol --- .../quantizer/fp32_quantizer/fp32_quantizer.cc | 3 +-- src/turbo/quantizer/quantizer.h | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index f5c29b76b..839137e6d 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -17,7 +17,6 @@ #include #include #include -#include #include #include #include "core/quantizer/record_quantizer.h" @@ -53,7 +52,7 @@ int Fp32Quantizer::init(const IndexMeta &meta, int Fp32Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, std::string *out, IndexQueryMeta *ometa) const { if (qmeta.unit_size() != sizeof(float)) { - return IndexError_Unsupported; + return kErrUnsupported; } // qmeta.dimension() may be the inflated (data + extras) dimension when the diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 3e5e65ec6..3b7a7d944 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -28,6 +28,20 @@ namespace turbo { using namespace zvec::core; +//! Error code literals mirroring core::IndexError::Code integer values. +//! +//! Turbo quantizer sources use these directly instead of the +//! `IndexError_NotImplemented` / `IndexError_Unsupported` const objects +//! because MSVC's WINDOWS_EXPORT_ALL_SYMBOLS does not export const data +//! with constructors from zvec_shared.dll. zvec_turbo is a static library +//! linked with /WHOLEARCHIVE, so referencing those unexported symbols across +//! the DLL boundary triggers LNK2019 on Windows. +//! +//! IndexError::Code stores -val in its constructor, so NotImplemented(11) +//! yields -11 and Unsupported(12) yields -12. +constexpr int kErrNotImplemented = -11; +constexpr int kErrUnsupported = -12; + //! Magic number ('QTZR') stamped at the start of a serialized quantizer blob. constexpr uint32_t kQuantizerMagic = 0x52545A51u; From e16cda561e2630b10bd805ccca3aa3c9e54c8089 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 2 Jul 2026 16:47:42 +0800 Subject: [PATCH 19/45] add untest --- tests/turbo/turbo_fht_rotator_test.cc | 255 ++++++++++++++++++++++++++ 1 file changed, 255 insertions(+) diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc index 178dad3fd..6acdaba80 100644 --- a/tests/turbo/turbo_fht_rotator_test.cc +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -202,3 +202,258 @@ TEST(FhtRotator, FromBlobMalformed) { hdr.payload_size = 0; EXPECT_EQ(FhtRotator::from_blob(&hdr, sizeof(hdr)), nullptr); } + +// --------------------------------------------------------------------------- +// L2 distance preserved (orthogonal transform) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, L2DistancePreserved) { + std::mt19937 gen(2024); + std::uniform_real_distribution dist(-1.0f, 1.0f); + + for (int dim : {32, 64, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + const int N = 50; + std::vector> raw(N, std::vector(dim)); + std::vector> rotated(N, std::vector(dim)); + for (int i = 0; i < N; ++i) { + fill_random(raw[i].data(), dim, gen); + rot->apply(raw[i].data(), rotated[i].data()); + } + + // Check that ||rotated[i] - rotated[j]|| ≈ ||raw[i] - raw[j]||. + for (int i = 1; i < N; ++i) { + float d_raw = 0.0f, d_rot = 0.0f; + for (int j = 0; j < dim; ++j) { + float dr = raw[i][j] - raw[0][j]; + float dt = rotated[i][j] - rotated[0][j]; + d_raw += dr * dr; + d_rot += dt * dt; + } + EXPECT_NEAR(d_raw, d_rot, 1e-2f) + << "L2 mismatch for dim=" << dim << " i=" << i; + } + } +} + +// --------------------------------------------------------------------------- +// Cosine distance preserved (orthogonal transform) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, CosineDistancePreserved) { + std::mt19937 gen(777); + std::uniform_real_distribution dist(0.1f, 1.0f); + + auto cosine_dist = [](const float *a, const float *b, int dim) { + float dot = 0, na = 0, nb = 0; + for (int i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom; + }; + + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + const int N = 50; + std::vector> raw(N, std::vector(dim)); + std::vector> rotated(N, std::vector(dim)); + for (int i = 0; i < N; ++i) { + fill_random(raw[i].data(), dim, gen); + rot->apply(raw[i].data(), rotated[i].data()); + } + + for (int i = 1; i < N; ++i) { + float d_raw = cosine_dist(raw[i].data(), raw[0].data(), dim); + float d_rot = cosine_dist(rotated[i].data(), rotated[0].data(), dim); + EXPECT_NEAR(d_raw, d_rot, 1e-3f) + << "Cosine mismatch for dim=" << dim << " i=" << i; + } + } +} + +// --------------------------------------------------------------------------- +// Apply is non-trivial (not identity) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, ApplyIsNonTrivial) { + std::mt19937 gen(42); + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector output(dim); + rot->apply(input.data(), output.data()); + + // At least some elements should differ from the input. + bool any_diff = false; + for (int i = 0; i < dim; ++i) { + if (std::abs(input[i] - output[i]) > 1e-6f) { + any_diff = true; + break; + } + } + EXPECT_TRUE(any_diff) << "apply is identity for dim=" << dim; + } +} + +// --------------------------------------------------------------------------- +// Apply is deterministic (same input → same output) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, ApplyDeterministic) { + std::mt19937 gen(55); + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector r1(dim), r2(dim); + rot->apply(input.data(), r1.data()); + rot->apply(input.data(), r2.data()); + + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) + << "non-deterministic apply at i=" << i << " dim=" << dim; + } + } +} + +// --------------------------------------------------------------------------- +// Deserialize on existing object (init → serialize → deserialize on new object) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, DeserializeOnExistingObject) { + std::mt19937 gen(314); + for (int dim : {32, 97, 128}) { + // Build and train original. + auto rot1 = FhtRotator::create(dim); + ASSERT_TRUE(rot1); + rot1->train(nullptr, 0, 0); + + std::string blob; + ASSERT_EQ(0, rot1->serialize(&blob)); + + // Create a fresh rotator, then call deserialize() on it. + auto rot2 = FhtRotator::create(dim); + ASSERT_TRUE(rot2); + ASSERT_EQ(0, rot2->deserialize(blob.data(), blob.size())); + + EXPECT_EQ(rot2->in_dim(), dim); + EXPECT_EQ(rot2->out_dim(), dim); + + // Both rotators should produce identical results. + std::vector input(dim); + fill_random(input.data(), dim, gen); + + std::vector r1(dim), r2(dim); + rot1->apply(input.data(), r1.data()); + rot2->apply(input.data(), r2.data()); + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) + << "apply mismatch at i=" << i << " dim=" << dim; + } + + // Inverse via rot2 should recover input. + check_round_trip(*rot2, input); + } +} + +// --------------------------------------------------------------------------- +// Deserialize with truncated payload fails +// --------------------------------------------------------------------------- + +TEST(FhtRotator, DeserializeTruncatedPayload) { + std::mt19937 gen(42); + auto rot = FhtRotator::create(64); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + std::string blob; + ASSERT_EQ(0, rot->serialize(&blob)); + + // Truncate the blob: keep header but cut half the payload. + auto *hdr = reinterpret_cast(blob.data()); + size_t truncated_len = sizeof(RotatorSerHeader) + hdr->payload_size / 2; + + auto rot2 = FhtRotator::create(64); + ASSERT_TRUE(rot2); + EXPECT_NE(0, rot2->deserialize(blob.data(), truncated_len)); + + // Also test from_blob with truncated data. + EXPECT_EQ(FhtRotator::from_blob(blob.data(), truncated_len), nullptr); +} + +// --------------------------------------------------------------------------- +// Large dimension stress test +// --------------------------------------------------------------------------- + +TEST(FhtRotator, LargeDimension) { + std::mt19937 gen(2025); + for (int dim : {1024, 2048, 4096}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot) << "create failed for dim=" << dim; + rot->train(nullptr, 0, 0); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + check_round_trip(*rot, input, 1e-2f); + + // Verify serialize/deserialize round-trip. + std::string blob; + ASSERT_EQ(0, rot->serialize(&blob)); + auto rot2 = FhtRotator::from_blob(blob.data(), blob.size()); + ASSERT_TRUE(rot2); + + std::vector r1(dim), r2(dim); + rot->apply(input.data(), r1.data()); + rot2->apply(input.data(), r2.data()); + for (int i = 0; i < dim; ++i) { + EXPECT_FLOAT_EQ(r1[i], r2[i]) + << "mismatch at i=" << i << " dim=" << dim; + } + } +} + +// --------------------------------------------------------------------------- +// Norm preserved (orthogonal transform preserves vector norm) +// --------------------------------------------------------------------------- + +TEST(FhtRotator, NormPreserved) { + std::mt19937 gen(99); + for (int dim : {32, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + rot->train(nullptr, 0, 0); + + std::vector input(dim); + fill_random(input.data(), dim, gen); + + float norm_in = 0.0f; + for (int i = 0; i < dim; ++i) norm_in += input[i] * input[i]; + + std::vector output(dim); + rot->apply(input.data(), output.data()); + + float norm_out = 0.0f; + for (int i = 0; i < dim; ++i) norm_out += output[i] * output[i]; + + EXPECT_NEAR(norm_in, norm_out, 1e-2f) + << "norm not preserved for dim=" << dim; + } +} From f3a67820ee31c4288614ff7d0ed3beaa9d9b7b60 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Thu, 2 Jul 2026 17:34:32 +0800 Subject: [PATCH 20/45] 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 79a08eb16ccd15e28f2b6404c32bcdf99684e945 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 11:15:01 +0800 Subject: [PATCH 21/45] 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 22/45] 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 23/45] 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 281fdb40b9fc72eba9aabdae939b3f6cdc15f20d Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 3 Jul 2026 14:16:23 +0800 Subject: [PATCH 24/45] 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 25/45] 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 8406e9fdc34e45e0990daba6f527faff0f7022d4 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Tue, 7 Jul 2026 20:30:43 +0800 Subject: [PATCH 26/45] 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 ac17ee916c1d0851e0b18a9d541b26fbf6888f87 Mon Sep 17 00:00:00 2001 From: ray Date: Fri, 10 Jul 2026 12:45:00 +0800 Subject: [PATCH 27/45] 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 28/45] 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 6abf1b8b25e6ccb9ff87343b20511eb9785859e7 Mon Sep 17 00:00:00 2001 From: zzl <1581199236@qq.com> Date: Fri, 10 Jul 2026 15:30:34 +0800 Subject: [PATCH 29/45] 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 30/45] 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 31/45] 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 32/45] 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 33/45] 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 34/45] 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 428713154e48ba7326f0091c4f0221096e2eb2f9 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 15 Jul 2026 10:33:01 +0800 Subject: [PATCH 35/45] 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 36/45] 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 084b475b4a81157bea84bb283f7e8890e80e2dd4 Mon Sep 17 00:00:00 2001 From: ray Date: Wed, 15 Jul 2026 16:44:01 +0800 Subject: [PATCH 37/45] 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 38/45] 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 39/45] 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 40/45] 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 41/45] 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 42/45] 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 43/45] 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 44/45] 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 45/45] 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}; } }