From 90012e87b4ea9e5b4f7a33bfc5bcf8e34f4e82cd Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Mon, 20 Jul 2026 16:19:58 +0800 Subject: [PATCH 1/2] feat(quantization): expose uniform uint7 and add uint8 Validate UniformUint7 dimensions and reformer inputs, and cover UniformUint8 graph and brute-force search through the Vamana query metric path. --- src/core/algorithm/vamana/vamana_streamer.cc | 2 +- src/core/interface/index.cc | 9 +- src/core/metric/metric_params.h | 12 +- ...int8_metric.cc => uniform_uint7_metric.cc} | 31 +- src/core/metric/uniform_uint8_metric.cc | 268 ++++++++++ src/core/quantizer/quantizer_params.h | 18 +- ...onverter.cc => uniform_uint7_converter.cc} | 127 +++-- ..._reformer.cc => uniform_uint7_reformer.cc} | 108 ++-- src/core/quantizer/uniform_uint8_converter.cc | 331 +++++++++++++ src/core/quantizer/uniform_uint8_reformer.cc | 231 +++++++++ src/include/zvec/core/interface/index_param.h | 6 +- src/include/zvec/turbo/turbo.h | 13 +- .../quantize.cc | 19 +- .../quantize.h | 14 +- .../squared_euclidean.cc | 33 +- .../squared_euclidean.h | 22 +- .../uniform_uint8/squared_euclidean.cc | 462 ++++++++++++++++++ .../uniform_uint8/squared_euclidean.h | 47 ++ src/turbo/turbo.cc | 34 +- .../algorithm/vamana/vamana_streamer_test.cc | 107 ++++ tests/core/interface/index_interface_test.cc | 113 +++++ tests/core/interface/quantizer_param_test.cc | 35 ++ .../core/metric/uniform_uint7_metric_test.cc | 76 +++ .../core/metric/uniform_uint8_metric_test.cc | 410 ++++++++++++++++ ...test.cc => uniform_uint7_reformer_test.cc} | 246 ++++++++-- .../quantizer/uniform_uint8_reformer_test.cc | 342 +++++++++++++ 26 files changed, 2902 insertions(+), 214 deletions(-) rename src/core/metric/{uniform_int8_metric.cc => uniform_uint7_metric.cc} (79%) create mode 100644 src/core/metric/uniform_uint8_metric.cc rename src/core/quantizer/{uniform_int8_converter.cc => uniform_uint7_converter.cc} (69%) rename src/core/quantizer/{uniform_int8_reformer.cc => uniform_uint7_reformer.cc} (66%) create mode 100644 src/core/quantizer/uniform_uint8_converter.cc create mode 100644 src/core/quantizer/uniform_uint8_reformer.cc rename src/turbo/distance/avx512_vnni/{uniform_int8 => uniform_uint7}/quantize.cc (81%) rename src/turbo/distance/avx512_vnni/{uniform_int8 => uniform_uint7}/quantize.h (63%) rename src/turbo/distance/avx512_vnni/{uniform_int8 => uniform_uint7}/squared_euclidean.cc (87%) rename src/turbo/distance/avx512_vnni/{uniform_int8 => uniform_uint7}/squared_euclidean.h (58%) create mode 100644 src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.cc create mode 100644 src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.h create mode 100644 tests/core/interface/quantizer_param_test.cc create mode 100644 tests/core/metric/uniform_uint7_metric_test.cc create mode 100644 tests/core/metric/uniform_uint8_metric_test.cc rename tests/core/quantizer/{uniform_int8_reformer_test.cc => uniform_uint7_reformer_test.cc} (65%) create mode 100644 tests/core/quantizer/uniform_uint8_reformer_test.cc diff --git a/src/core/algorithm/vamana/vamana_streamer.cc b/src/core/algorithm/vamana/vamana_streamer.cc index d651c486b..e81d191bb 100644 --- a/src/core/algorithm/vamana/vamana_streamer.cc +++ b/src/core/algorithm/vamana/vamana_streamer.cc @@ -236,7 +236,7 @@ int VamanaStreamer::open(IndexStorage::Pointer stg) { metric_params.merge(meta_.metric_params()); meta_.set_metric(index_meta.metric_name(), 0, metric_params); // Propagate reformer info from stored meta (needed for quantizers - // whose reformer params are computed during training, e.g. UniformInt8) + // whose reformer params are computed during training, e.g. UniformUint7) if (!index_meta.reformer_name().empty()) { meta_.set_reformer(index_meta.reformer_name(), 0, index_meta.reformer_params()); diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 73dd7950d..f76c321ca 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -212,8 +212,11 @@ int Index::CreateAndInitConverterReformer(const QuantizerParam ¶m, case QuantizerType::kRabitq: // no converter here return 0; - case QuantizerType::kUniformInt8: - converter_name = "UniformInt8StreamingConverter"; + case QuantizerType::kUniformUint7: + converter_name = "UniformUint7StreamingConverter"; + break; + case QuantizerType::kUniformUint8: + converter_name = "UniformUint8StreamingConverter"; break; default: LOG_ERROR("Unsupported quantizer type: "); @@ -371,7 +374,7 @@ int Index::Open(const std::string &file_path, StorageOptions storage_options) { } // If a converter exists but reformer was not created during Init() - // (converters like UniformInt8 whose reformer params are only available + // (uniform quantizers whose reformer params are only available // after train()), create it now from the persisted meta that the streamer // has loaded. When there is no converter (QuantizerType::kNone), reformer_ // is nullptr by design — skip this block entirely. diff --git a/src/core/metric/metric_params.h b/src/core/metric/metric_params.h index 05f8db96a..824e31a33 100644 --- a/src/core/metric/metric_params.h +++ b/src/core/metric/metric_params.h @@ -34,9 +34,13 @@ static const std::string QUANTIZED_INTEGER_METRIC_ORIGIN_METRIC_NAME = static const std::string QUANTIZED_INTEGER_METRIC_ORIGIN_METRIC_PARAMS = "proxima.quantized_integer.metric.origin_metric_params"; -//! UniformInt8 Metric -static const std::string UNIFORM_INT8_METRIC_ORIGIN_METRIC_NAME = - "proxima.uniform_int8.metric.origin_metric_name"; +//! UniformUint7 Metric +static const std::string UNIFORM_UINT7_METRIC_ORIGIN_METRIC_NAME = + "proxima.uniform_uint7.metric.origin_metric_name"; + +//! UniformUint8 Metric +static const std::string UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME = + "proxima.uniform_uint8.metric.origin_metric_name"; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/metric/uniform_int8_metric.cc b/src/core/metric/uniform_uint7_metric.cc similarity index 79% rename from src/core/metric/uniform_int8_metric.cc rename to src/core/metric/uniform_uint7_metric.cc index 32fb410c6..852f1e283 100644 --- a/src/core/metric/uniform_int8_metric.cc +++ b/src/core/metric/uniform_uint7_metric.cc @@ -16,38 +16,44 @@ #include #include #include +#include #include #include "metric_params.h" namespace zvec { namespace core { -/*! Index Metric for Uniform Int8 Quantization (Global Scale) +/*! Index Metric for Uniform Uint7 Quantization (Global Scale) * * Uses direct int8 L2 distance computation. Since all vectors share * a single global scale/bias, no per-vector reconstruction is needed. * This is the key benefit: distance = sum((a[i] - b[i])^2) on raw int8 * values, with optional post-scaling by 1/scale^2 for real L2 distances. */ -class UniformInt8Metric : public IndexMetric { +class UniformUint7Metric : public IndexMetric { public: //! Initialize Metric int init(const IndexMeta &meta, const ailego::Params &index_params) override { if (meta.data_type() != IndexMeta::DataType::DT_INT8) { - LOG_ERROR("UniformInt8Metric: unsupported type %d", meta.data_type()); + LOG_ERROR("UniformUint7Metric: unsupported type %d", meta.data_type()); return IndexError_Unsupported; } + if (meta.dimension() == 0 || meta.dimension() > MAX_DIMENSION) { + LOG_ERROR("UniformUint7Metric: dimension=%u must be in [1, %d]", + meta.dimension(), MAX_DIMENSION); + return IndexError_InvalidArgument; + } std::string metric_name; - index_params.get(UNIFORM_INT8_METRIC_ORIGIN_METRIC_NAME, &metric_name); + index_params.get(UNIFORM_UINT7_METRIC_ORIGIN_METRIC_NAME, &metric_name); if (metric_name.empty()) { - LOG_ERROR("UniformInt8Metric: param %s is required", - UNIFORM_INT8_METRIC_ORIGIN_METRIC_NAME.c_str()); + LOG_ERROR("UniformUint7Metric: param %s is required", + UNIFORM_UINT7_METRIC_ORIGIN_METRIC_NAME.c_str()); return IndexError_InvalidArgument; } if (metric_name != "SquaredEuclidean") { - LOG_ERROR("UniformInt8Metric: only SquaredEuclidean supported, got %s", + LOG_ERROR("UniformUint7Metric: only SquaredEuclidean supported, got %s", metric_name.c_str()); return IndexError_Unsupported; } @@ -55,7 +61,7 @@ class UniformInt8Metric : public IndexMetric { meta_ = meta; params_ = index_params; - LOG_INFO("UniformInt8Metric initialized: dimension=%u", meta_.dimension()); + LOG_INFO("UniformUint7Metric initialized: dimension=%u", meta_.dimension()); return 0; } @@ -67,15 +73,16 @@ class UniformInt8Metric : public IndexMetric { //! Retrieve if it matched bool is_matched(const IndexMeta &meta) const override { return meta.data_type() == meta_.data_type() && - meta.unit_size() == meta_.unit_size(); + meta.unit_size() == meta_.unit_size() && + meta.dimension() == meta_.dimension(); } //! Retrieve if it matched bool is_matched(const IndexMeta &meta, const IndexQueryMeta &qmeta) const override { - return qmeta.data_type() == meta_.data_type() && + return is_matched(meta) && qmeta.data_type() == meta_.data_type() && qmeta.unit_size() == meta_.unit_size() && - qmeta.dimension() == meta.dimension(); + qmeta.dimension() == meta_.dimension(); } //! Retrieve distance function for query (1x1) @@ -152,7 +159,7 @@ class UniformInt8Metric : public IndexMetric { ailego::Params params_{}; }; -INDEX_FACTORY_REGISTER_METRIC_ALIAS(UniformInt8, UniformInt8Metric); +INDEX_FACTORY_REGISTER_METRIC_ALIAS(UniformUint7, UniformUint7Metric); } // namespace core } // namespace zvec diff --git a/src/core/metric/uniform_uint8_metric.cc b/src/core/metric/uniform_uint8_metric.cc new file mode 100644 index 000000000..68d5ddbf7 --- /dev/null +++ b/src/core/metric/uniform_uint8_metric.cc @@ -0,0 +1,268 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include "metric_params.h" + +namespace zvec { +namespace core { + +namespace { + +constexpr size_t kTailBytes = sizeof(uint32_t); + +size_t OriginalDimension(size_t encoded_dimension) { + return encoded_dimension > kTailBytes ? encoded_dimension - kTailBytes : 0; +} + +void UniformUint8QueryPreprocess(void *query, size_t encoded_dimension) { + const size_t original_dimension = OriginalDimension(encoded_dimension); + if (original_dimension == 0) { + return; + } + + const auto *raw_query = static_cast(query); + // The tail may already contain a correction from an earlier call. Derive + // everything from the immutable query body so preprocessing is idempotent. + uint64_t sum = 0; + uint64_t sum_squared = 0; + for (size_t i = 0; i < original_dimension; ++i) { + const uint64_t value = raw_query[i]; + sum += value; + sum_squared += value * value; + } + + const int64_t correction = + static_cast(sum_squared) - 256 * static_cast(sum); + if (correction < (std::numeric_limits::min)() || + correction > (std::numeric_limits::max)()) { + // Oversized externally supplied metadata uses the scalar squared- + // difference path, which ignores the query tail. + return; + } + const int32_t encoded_correction = static_cast(correction); + std::memcpy(static_cast(query) + original_dimension, + &encoded_correction, sizeof(encoded_correction)); +} + +IndexMetric::DistanceBatchQueryPreprocessFunc +UniformUint8QueryPreprocessFunc() { + static const auto preprocess = []() { + auto turbo_preprocess = turbo::get_query_preprocess_func( + turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8, + turbo::QuantizeType::kUniformUint8); + return turbo_preprocess ? turbo_preprocess : UniformUint8QueryPreprocess; + }(); + return preprocess; +} + +void UniformUint8StoredSquaredEuclidean(const void *lhs_data, + const void *rhs_data, + size_t encoded_dimension, + float *distance) { + const size_t original_dimension = OriginalDimension(encoded_dimension); + const auto *lhs = static_cast(lhs_data); + const auto *rhs = static_cast(rhs_data); + int64_t sum = 0; + for (size_t i = 0; i < original_dimension; ++i) { + const int difference = static_cast(lhs[i]) - static_cast(rhs[i]); + sum += static_cast(difference) * difference; + } + *distance = static_cast(sum); +} + +IndexMetric::MatrixDistance UniformUint8StoredDistance() { + static const IndexMetric::MatrixDistance distance = []() { + auto turbo_distance = turbo::get_distance_func( + turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8, + turbo::QuantizeType::kUniformUint8); + return turbo_distance ? turbo_distance : UniformUint8StoredSquaredEuclidean; + }(); + return distance; +} + +void UniformUint8StoredSquaredEuclideanBatch(const void *const *vectors, + const void *query, size_t count, + size_t encoded_dimension, + float *distances) { + const auto distance = UniformUint8StoredDistance(); + for (size_t i = 0; i < count; ++i) { + distance(vectors[i], query, encoded_dimension, distances + i); + } +} + +void UniformUint8StoredQuerySquaredEuclidean(const void *stored_data, + const void *query_data, + size_t encoded_dimension, + float *distance) { + const size_t original_dimension = OriginalDimension(encoded_dimension); + const auto *stored = static_cast(stored_data); + const auto *query = static_cast(query_data); + int64_t sum = 0; + for (size_t i = 0; i < original_dimension; ++i) { + const int difference = + static_cast(stored[i]) - (static_cast(query[i]) - 128); + sum += static_cast(difference) * difference; + } + *distance = static_cast(sum); +} + +void UniformUint8StoredQuerySquaredEuclideanBatch(const void *const *vectors, + const void *query, + size_t count, + size_t encoded_dimension, + float *distances) { + for (size_t i = 0; i < count; ++i) { + UniformUint8StoredQuerySquaredEuclidean(vectors[i], query, + encoded_dimension, distances + i); + } +} + +} // namespace + +class UniformUint8QueryMetric : public IndexMetric { + public: + UniformUint8QueryMetric() = default; + UniformUint8QueryMetric(const IndexMeta &meta, const ailego::Params ¶ms) + : meta_(meta), params_(params) {} + + int init(const IndexMeta &meta, const ailego::Params ¶ms) override { + if (meta.data_type() != IndexMeta::DataType::DT_INT8) { + LOG_ERROR("UniformUint8Metric: unsupported type %d", meta.data_type()); + return IndexError_Unsupported; + } + if (meta.dimension() <= kTailBytes) { + LOG_ERROR( + "UniformUint8Metric: encoded dimension=%u must include a non-empty " + "vector and a %zu-byte tail", + meta.dimension(), kTailBytes); + return IndexError_InvalidArgument; + } + + std::string metric_name; + params.get(UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME, &metric_name); + if (metric_name.empty()) { + LOG_ERROR("UniformUint8Metric: param %s is required", + UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME.c_str()); + return IndexError_InvalidArgument; + } + if (metric_name != "SquaredEuclidean") { + LOG_ERROR("UniformUint8Metric: only SquaredEuclidean supported, got %s", + metric_name.c_str()); + return IndexError_Unsupported; + } + + meta_ = meta; + params_ = params; + return 0; + } + + int cleanup(void) override { + return 0; + } + + bool is_matched(const IndexMeta &meta) const override { + return meta.data_type() == meta_.data_type() && + meta.unit_size() == meta_.unit_size() && + meta.dimension() == meta_.dimension(); + } + + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &query_meta) const override { + return is_matched(meta) && query_meta.data_type() == meta_.data_type() && + query_meta.unit_size() == meta_.unit_size() && + query_meta.dimension() == meta_.dimension(); + } + + MatrixDistance distance(void) const override { + return UniformUint8StoredQuerySquaredEuclidean; + } + + MatrixBatchDistance batch_distance(void) const override { + const size_t original_dimension = OriginalDimension(meta_.dimension()); + // The VNNI kernel reduces its signed dot product in int32 lanes. The + // public quantizer dimension bound guarantees that reduction is exact; + // preserve the scalar int64 path for larger externally supplied metadata. + if (original_dimension > 0 && original_dimension <= MAX_DIMENSION) { + auto turbo_distance = turbo::get_batch_distance_func( + turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8, + turbo::QuantizeType::kUniformUint8); + if (turbo_distance) { + return turbo_distance; + } + } + return UniformUint8StoredQuerySquaredEuclideanBatch; + } + + DistanceBatchQueryPreprocessFunc get_query_preprocess_func() const override { + return UniformUint8QueryPreprocessFunc(); + } + + const ailego::Params ¶ms(void) const override { + return params_; + } + + int train(const void * /*vector*/, size_t /*dimension*/) override { + return 0; + } + + bool support_train(void) const override { + return false; + } + + void normalize(float * /*score*/) const override {} + + bool support_normalize(void) const override { + return false; + } + + Pointer query_metric(void) const override { + return nullptr; + } + + protected: + IndexMeta meta_{}; + ailego::Params params_{}; +}; + +class UniformUint8Metric : public UniformUint8QueryMetric { + public: + MatrixDistance distance(void) const override { + return UniformUint8StoredDistance(); + } + + MatrixDistance distance_matrix(size_t rows, size_t columns) const override { + return rows == 1 && columns == 1 ? UniformUint8StoredDistance() : nullptr; + } + + MatrixBatchDistance batch_distance(void) const override { + return UniformUint8StoredSquaredEuclideanBatch; + } + + Pointer query_metric(void) const override { + return std::make_shared(meta_, params_); + } +}; + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(UniformUint8, UniformUint8Metric); + +} // namespace core +} // namespace zvec diff --git a/src/core/quantizer/quantizer_params.h b/src/core/quantizer/quantizer_params.h index d56c8591d..ab04d462a 100644 --- a/src/core/quantizer/quantizer_params.h +++ b/src/core/quantizer/quantizer_params.h @@ -119,11 +119,17 @@ static const std::string INTEGER_STREAMING_REFORMER_ENABLE_NORMALIZE = static const std::string INTEGER_STREAMING_REFORMER_IS_EUCLIDEAN = "integer_streaming.reformer.is_euclidean"; -//! UniformInt8StreamingConverter / Reformer -static const std::string UNIFORM_INT8_REFORMER_SCALE = - "uniform_int8.reformer.scale"; -static const std::string UNIFORM_INT8_REFORMER_BIAS = - "uniform_int8.reformer.bias"; +//! UniformUint7StreamingConverter / Reformer +static const std::string UNIFORM_UINT7_REFORMER_SCALE = + "uniform_uint7.reformer.scale"; +static const std::string UNIFORM_UINT7_REFORMER_BIAS = + "uniform_uint7.reformer.bias"; + +//! UniformUint8StreamingConverter / Reformer +static const std::string UNIFORM_UINT8_REFORMER_SCALE = + "uniform_uint8.reformer.scale"; +static const std::string UNIFORM_UINT8_REFORMER_BIAS = + "uniform_uint8.reformer.bias"; //! DoubleBitConverter static const std::string DOUBLE_BIT_CONVERTER_TRAIN_SAMPLE_COUNT = @@ -140,4 +146,4 @@ static const std::string DOUBLE_BIT_REFORMER_B_VALUE = "double_bit.reformer.b_value"; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/quantizer/uniform_int8_converter.cc b/src/core/quantizer/uniform_uint7_converter.cc similarity index 69% rename from src/core/quantizer/uniform_int8_converter.cc rename to src/core/quantizer/uniform_uint7_converter.cc index 5b23210de..3621f13cc 100644 --- a/src/core/quantizer/uniform_int8_converter.cc +++ b/src/core/quantizer/uniform_uint7_converter.cc @@ -19,67 +19,77 @@ #include #include #include +#include #include #include "../metric/metric_params.h" namespace zvec { namespace core { -/*! Converter for Uniform Int8 Quantization (Global Scale) +/*! Converter for Uniform Uint7 Quantization (Global Scale) * * Unlike IntegerStreamingConverter which uses per-vector scale/bias, * this converter computes a single global scale/bias from the entire dataset. * All vectors share the same quantization parameters, enabling direct int8 * L2 distance computation without per-vector reconstruction. */ -class UniformInt8StreamingConverter : public IndexConverter { +class UniformUint7StreamingConverter : public IndexConverter { public: //! Constructor. //! `dst_type` is required by the INDEX_FACTORY_REGISTER_CONVERTER_ALIAS //! macro signature but is unused here: the output type is always //! IndexMeta::DataType::DT_INT8, hard-coded in init(). - UniformInt8StreamingConverter(IndexMeta::DataType /*dst_type*/) {} + UniformUint7StreamingConverter(IndexMeta::DataType /*dst_type*/) {} //! Destructor - ~UniformInt8StreamingConverter() override {} + ~UniformUint7StreamingConverter() override {} //! Initialize Converter int init(const IndexMeta &index_meta, const ailego::Params ¶ms) override { meta_ = index_meta; original_dimension_ = index_meta.dimension(); + scale_ = 0.0f; + bias_ = 0.0f; + holder_.reset(); // Reset stats so a re-init() call does not leak counters from a // previous lifecycle. *stats_.mutable_trained_count() = 0; *stats_.mutable_transformed_count() = 0; + if (original_dimension_ == 0 || original_dimension_ > MAX_DIMENSION) { + LOG_ERROR( + "UniformUint7StreamingConverter: dimension=%zu must be in [1, %d]", + original_dimension_, MAX_DIMENSION); + return IndexError_InvalidArgument; + } + // Store converter info in meta - meta_.set_converter("UniformInt8StreamingConverter", 0, params); + meta_.set_converter("UniformUint7StreamingConverter", 0, params); // Set data type to INT8, dimension stays the same (no per-vector extras) meta_.set_meta(IndexMeta::DataType::DT_INT8, original_dimension_); // Set metric to our direct int8 L2 metric ailego::Params metric_params; - metric_params.set(UNIFORM_INT8_METRIC_ORIGIN_METRIC_NAME, + metric_params.set(UNIFORM_UINT7_METRIC_ORIGIN_METRIC_NAME, index_meta.metric_name()); - meta_.set_metric("UniformInt8", 0, metric_params); - - // Restore scale/bias from persisted params if available (search-only - // path where train() is skipped). Otherwise they stay at 0 and will - // be computed in train(). - params.get(UNIFORM_INT8_REFORMER_SCALE, &scale_); - params.get(UNIFORM_INT8_REFORMER_BIAS, &bias_); - - // Only register reformer when scale/bias are available (either from - // persisted params or after train()). During first-time Init() before - // train(), we intentionally leave reformer_name empty so that the - // Index layer does not attempt to create an uninitialized reformer. - if (scale_ != 0.0f) { + meta_.set_metric("UniformUint7", 0, metric_params); + + const bool has_scale = params.get(UNIFORM_UINT7_REFORMER_SCALE, &scale_); + const bool has_bias = params.get(UNIFORM_UINT7_REFORMER_BIAS, &bias_); + if (has_scale != has_bias || + (has_scale && + (!std::isfinite(scale_) || scale_ <= 0.0f || !std::isfinite(bias_)))) { + LOG_ERROR("UniformUint7StreamingConverter: invalid scale/bias params"); + return IndexError_InvalidArgument; + } + + if (has_scale) { ailego::Params reformer_params; - reformer_params.set(UNIFORM_INT8_REFORMER_SCALE, scale_); - reformer_params.set(UNIFORM_INT8_REFORMER_BIAS, bias_); - meta_.set_reformer("UniformInt8StreamingReformer", 0, reformer_params); + reformer_params.set(UNIFORM_UINT7_REFORMER_SCALE, scale_); + reformer_params.set(UNIFORM_UINT7_REFORMER_BIAS, bias_); + meta_.set_reformer("UniformUint7StreamingReformer", 0, reformer_params); } return 0; @@ -89,25 +99,33 @@ class UniformInt8StreamingConverter : public IndexConverter { int cleanup(void) override { *stats_.mutable_trained_count() = 0; *stats_.mutable_transformed_count() = 0; + holder_.reset(); + scale_ = 0.0f; + bias_ = 0.0f; return 0; } //! Train: compute global min/max and derive scale/bias int train(IndexHolder::Pointer holder) override { if (!holder) { - LOG_ERROR("UniformInt8StreamingConverter: null holder in train"); + LOG_ERROR("UniformUint7StreamingConverter: null holder in train"); return IndexError_InvalidArgument; } + if (holder->data_type() != IndexMeta::DataType::DT_FP32 || + holder->dimension() != original_dimension_) { + return IndexError_Mismatch; + } ailego::ElapsedTime timer; AILEGO_DEFER([&]() { stats_.set_trained_costtime(timer.milli_seconds()); }); + *stats_.mutable_trained_count() = 0; float global_min = std::numeric_limits::max(); float global_max = std::numeric_limits::lowest(); auto iter = holder->create_iterator(); if (!iter) { - LOG_ERROR("UniformInt8StreamingConverter: failed to create iterator"); + LOG_ERROR("UniformUint7StreamingConverter: failed to create iterator"); return IndexError_Runtime; } @@ -118,7 +136,7 @@ class UniformInt8StreamingConverter : public IndexConverter { float v = vec[i]; if (!std::isfinite(v)) { LOG_ERROR( - "UniformInt8StreamingConverter: non-finite value in training " + "UniformUint7StreamingConverter: non-finite value in training " "set (record_idx=%zu, dim_idx=%zu, value=%f)", (size_t)*stats_.mutable_trained_count(), i, v); return IndexError_InvalidArgument; @@ -135,12 +153,13 @@ class UniformInt8StreamingConverter : public IndexConverter { // Reject empty training set: scale/bias would be undefined and would // silently produce all-clipped int8 vectors at search time. if (*stats_.mutable_trained_count() == 0) { - LOG_ERROR("UniformInt8StreamingConverter: empty training set"); + LOG_ERROR("UniformUint7StreamingConverter: empty training set"); return IndexError_InvalidArgument; } // Compute global scale and bias: - // forward: int8 = clip(round(float * scale + bias), 0, 127) + // forward: + // int8 = clip(round-to-nearest-even(float * scale + bias), 0, 127) // inverse: float ≈ (int8 - bias) / scale // // Values are mapped to [0, 127] to enable the VNNI abs trick in the @@ -150,7 +169,11 @@ class UniformInt8StreamingConverter : public IndexConverter { // Lossless integer fast-path: when all training values are integers and // the dynamic range fits within 127, we use scale=1 for exact mapping. constexpr float epsilon = std::numeric_limits::epsilon(); - float range = global_max - global_min; + const float range = global_max - global_min; + if (!std::isfinite(range)) { + LOG_ERROR("UniformUint7StreamingConverter: non-finite training range"); + return IndexError_InvalidArgument; + } if (all_integer && range <= 127.0f) { scale_ = 1.0f; bias_ = -global_min; // global_min is integer — maps to 0 @@ -160,35 +183,37 @@ class UniformInt8StreamingConverter : public IndexConverter { } LOG_INFO( - "UniformInt8StreamingConverter train done: costtime %zums, " + "UniformUint7StreamingConverter train done: costtime %zums, " "global_min=%f, global_max=%f, scale=%f, bias=%f", (size_t)timer.milli_seconds(), global_min, global_max, scale_, bias_); // Now configure the reformer with the computed scale/bias ailego::Params reformer_params; - reformer_params.set(UNIFORM_INT8_REFORMER_SCALE, scale_); - reformer_params.set(UNIFORM_INT8_REFORMER_BIAS, bias_); - meta_.set_reformer("UniformInt8StreamingReformer", 0, reformer_params); + reformer_params.set(UNIFORM_UINT7_REFORMER_SCALE, scale_); + reformer_params.set(UNIFORM_UINT7_REFORMER_BIAS, bias_); + meta_.set_reformer("UniformUint7StreamingReformer", 0, reformer_params); - // Also store scale/bias in converter params for persistence + // Publish scale/bias through converter params so IndexMeta can carry them. ailego::Params conv_params = meta_.converter_params(); - conv_params.set(UNIFORM_INT8_REFORMER_SCALE, scale_); - conv_params.set(UNIFORM_INT8_REFORMER_BIAS, bias_); - meta_.set_converter(meta_.converter_name(), 0, conv_params); + conv_params.set(UNIFORM_UINT7_REFORMER_SCALE, scale_); + conv_params.set(UNIFORM_UINT7_REFORMER_BIAS, bias_); + meta_.set_converter(meta_.converter_name(), meta_.converter_revision(), + conv_params); return 0; } //! Transform: wrap holder to produce quantized int8 data int transform(IndexHolder::Pointer holder) override { - if (holder->data_type() != IndexMeta::DataType::DT_FP32 || - holder->dimension() != original_dimension_) { + if (!holder || holder->data_type() != IndexMeta::DataType::DT_FP32 || + holder->dimension() != original_dimension_ || !std::isfinite(scale_) || + scale_ <= 0.0f || !std::isfinite(bias_)) { return IndexError_Mismatch; } *stats_.mutable_transformed_count() += holder->count(); - holder_ = std::make_shared(holder, original_dimension_, - scale_, bias_); + holder_ = std::make_shared(holder, original_dimension_, + scale_, bias_); return 0; } @@ -214,12 +239,12 @@ class UniformInt8StreamingConverter : public IndexConverter { } private: - //! IndexHolder that applies uniform int8 quantization on-the-fly - class UniformInt8Holder : public IndexHolder { + //! IndexHolder that applies uniform uint7 quantization on-the-fly + class UniformUint7Holder : public IndexHolder { public: class Iterator : public IndexHolder::Iterator { public: - Iterator(const UniformInt8Holder *owner, + Iterator(const UniformUint7Holder *owner, IndexHolder::Iterator::Pointer &&iter) : owner_(owner), buffer_(owner->dimension(), 0), @@ -262,19 +287,21 @@ class UniformInt8StreamingConverter : public IndexConverter { return; } for (size_t i = 0; i < dim; ++i) { - float v = std::round(vec[i] * scale + bias); + // Match _mm512_cvtps_epi32 used by the SIMD kernel under the + // default floating-point rounding mode. + float v = std::nearbyint(vec[i] * scale + bias); v = std::max(0.0f, std::min(127.0f, v)); out[i] = static_cast(v); } } - const UniformInt8Holder *owner_{nullptr}; + const UniformUint7Holder *owner_{nullptr}; std::vector buffer_{}; IndexHolder::Iterator::Pointer front_iter_{}; }; - UniformInt8Holder(IndexHolder::Pointer front, size_t original_dim, - float scale, float bias) + UniformUint7Holder(IndexHolder::Pointer front, size_t original_dim, + float scale, float bias) : front_(std::move(front)), original_dim_(original_dim), scale_(scale), @@ -306,7 +333,7 @@ class UniformInt8StreamingConverter : public IndexConverter { IndexHolder::Iterator::Pointer create_iterator(void) override { auto iter = front_->create_iterator(); return iter ? IndexHolder::Iterator::Pointer( - new UniformInt8Holder::Iterator(this, std::move(iter))) + new UniformUint7Holder::Iterator(this, std::move(iter))) : IndexHolder::Iterator::Pointer(); } @@ -328,8 +355,8 @@ class UniformInt8StreamingConverter : public IndexConverter { float bias_{0.0f}; }; -INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(UniformInt8StreamingConverter, - UniformInt8StreamingConverter, +INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(UniformUint7StreamingConverter, + UniformUint7StreamingConverter, IndexMeta::DataType::DT_INT8); } // namespace core diff --git a/src/core/quantizer/uniform_int8_reformer.cc b/src/core/quantizer/uniform_uint7_reformer.cc similarity index 66% rename from src/core/quantizer/uniform_int8_reformer.cc rename to src/core/quantizer/uniform_uint7_reformer.cc index b642baf14..f8d239db8 100644 --- a/src/core/quantizer/uniform_int8_reformer.cc +++ b/src/core/quantizer/uniform_uint7_reformer.cc @@ -14,71 +14,67 @@ #include #include +#include #include #include +#include #include namespace zvec { namespace core { -/*! Reformer for Uniform Int8 Quantization (Global Scale) +/*! Reformer for Uniform Uint7 Quantization (Global Scale) * - * Uses a global scale/bias (computed by UniformInt8StreamingConverter) to + * Uses a global scale/bias (computed by UniformUint7StreamingConverter) to * quantize query vectors and build-time record vectors to int8. * No per-vector extras are appended — the output is pure int8. */ -class UniformInt8StreamingReformer : public IndexReformer { +class UniformUint7StreamingReformer : public IndexReformer { public: //! Constructor. //! `dst_type` is required by the INDEX_FACTORY_REGISTER_REFORMER_ALIAS //! macro signature but is unused here: the quantization output is //! always int8, governed by the (scale, bias) pair received in init(). - UniformInt8StreamingReformer(IndexMeta::DataType /*dst_type*/) {} + UniformUint7StreamingReformer(IndexMeta::DataType /*dst_type*/) {} //! Initialize Reformer //! - //! Lifecycle note: during build, scale/bias come from the converter's - //! train(); during search-only path, the converter first creates the - //! reformer with empty params, then Index::Open re-invokes init() with - //! the persisted params. We treat empty-params as "not yet initialized" - //! and reject any quantize/normalize call until real params arrive, so a - //! mis-wired pipeline fails loudly instead of silently producing garbage. + //! Scale/bias are supplied through reformer params after converter training. int init(const ailego::Params ¶ms) override { - bool has_scale = params.get(UNIFORM_INT8_REFORMER_SCALE, &scale_); - bool has_bias = params.get(UNIFORM_INT8_REFORMER_BIAS, &bias_); + Reset(); + float scale = 0.0f; + float bias = 0.0f; + const bool has_scale = params.get(UNIFORM_UINT7_REFORMER_SCALE, &scale); + const bool has_bias = params.get(UNIFORM_UINT7_REFORMER_BIAS, &bias); if (!has_scale || !has_bias) { LOG_ERROR( - "UniformInt8StreamingReformer init: missing required params " + "UniformUint7StreamingReformer: missing scale/bias params " "(scale_present=%d, bias_present=%d)", - (int)has_scale, (int)has_bias); - initialized_ = false; + static_cast(has_scale), static_cast(has_bias)); return IndexError_InvalidArgument; } - - if (!std::isfinite(scale_) || scale_ == 0.0f || !std::isfinite(bias_)) { + if (!std::isfinite(scale) || scale <= 0.0f || !std::isfinite(bias)) { LOG_ERROR( - "UniformInt8StreamingReformer: invalid params scale=%f, bias=%f", - scale_, bias_); - initialized_ = false; + "UniformUint7StreamingReformer: invalid params scale=%f, bias=%f", + scale, bias); return IndexError_InvalidArgument; } - // int8_l2 = scale^2 * real_l2, so real_l2 = int8_l2 / scale^2. + scale_ = scale; + bias_ = bias; scale_reciprocal_sq_ = 1.0f / (scale_ * scale_); initialized_ = true; - - // Resolve the SIMD quantize kernel once; falls back to scalar when the - // current CPU lacks AVX-512 (turbo returns nullptr on those builds). quantize_func_ = turbo::get_uniform_quantize_func(turbo::DataType::kInt8); - LOG_INFO("UniformInt8StreamingReformer init: scale=%f, bias=%f, simd=%s", + LOG_INFO("UniformUint7StreamingReformer init: scale=%f, bias=%f, simd=%s", scale_, bias_, quantize_func_ != nullptr ? "avx512" : "scalar"); return 0; } //! Cleanup Reformer int cleanup(void) override { + Reset(); return 0; } @@ -121,7 +117,7 @@ class UniformInt8StreamingReformer : public IndexReformer { IndexDocumentList &result) const override { if (!initialized_) { LOG_ERROR( - "UniformInt8StreamingReformer::normalize called before init " + "UniformUint7StreamingReformer::normalize called before init " "with valid params"); return IndexError_Runtime; } @@ -141,51 +137,81 @@ class UniformInt8StreamingReformer : public IndexReformer { std::string *out) const override { if (!initialized_) { LOG_ERROR( - "UniformInt8StreamingReformer::revert called before init " + "UniformUint7StreamingReformer::revert called before init " "with valid params"); return IndexError_Runtime; } - size_t dim = qmeta.dimension(); + if (!in || !out || qmeta.data_type() != IndexMeta::DataType::DT_INT8 || + qmeta.unit_size() != + IndexMeta::UnitSizeof(IndexMeta::DataType::DT_INT8)) { + return IndexError_InvalidArgument; + } + const size_t dim = qmeta.dimension(); + if (dim == 0 || dim > MAX_DIMENSION) { + LOG_ERROR( + "UniformUint7StreamingReformer: dimension=%zu must be in [1, %d]", + dim, MAX_DIMENSION); + return IndexError_InvalidArgument; + } out->resize(dim * sizeof(float)); - float *out_buf = reinterpret_cast(out->data()); - const int8_t *buf = reinterpret_cast(in); + char *out_buf = out->data(); + const auto *buf = static_cast(in); // Approximate dequantization (lossy): - // forward: int8 = clip(round(float * scale + bias), -127, 127) + // forward: + // int8 = clip(round-to-nearest-even(float * scale + bias), 0, 127) // inverse: float ≈ (int8 - bias) / scale // initialized_ guarantees scale_ != 0 and finite. float inv_scale = 1.0f / scale_; for (size_t i = 0; i < dim; ++i) { - out_buf[i] = (static_cast(buf[i]) - bias_) * inv_scale; + const float value = (static_cast(buf[i]) - bias_) * inv_scale; + std::memcpy(out_buf + i * sizeof(value), &value, sizeof(value)); } return 0; } private: + void Reset() { + scale_ = 0.0f; + bias_ = 0.0f; + scale_reciprocal_sq_ = 1.0f; + initialized_ = false; + quantize_func_ = nullptr; + } + //! Common quantization path shared by transform()/convert() (single & batch) int do_quantize(const void *src, const IndexQueryMeta &smeta, uint32_t count, std::string *out, IndexQueryMeta *ometa) const { if (!initialized_) { LOG_ERROR( - "UniformInt8StreamingReformer: quantize called before init " + "UniformUint7StreamingReformer: quantize called before init " "with valid params"); return IndexError_Runtime; } + if (!src || !out || !ometa || count == 0) { + return IndexError_InvalidArgument; + } if (smeta.data_type() != IndexMeta::DataType::DT_FP32 || smeta.unit_size() != IndexMeta::UnitSizeof(IndexMeta::DataType::DT_FP32)) { return IndexError_Unsupported; } + const size_t dim = smeta.dimension(); + if (dim == 0 || dim > MAX_DIMENSION) { + LOG_ERROR( + "UniformUint7StreamingReformer: dimension=%zu must be in [1, %d]", + dim, MAX_DIMENSION); + return IndexError_InvalidArgument; + } *ometa = smeta; - ometa->set_meta(IndexMeta::DataType::DT_INT8, smeta.dimension()); + ometa->set_meta(IndexMeta::DataType::DT_INT8, dim); const size_t out_stride = ometa->element_size(); out->resize(static_cast(count) * out_stride); - const float *vec = reinterpret_cast(src); - int8_t *ovec = reinterpret_cast(&(*out)[0]); - const size_t dim = smeta.dimension(); + const auto *vec = static_cast(src); + auto *ovec = reinterpret_cast(out->data()); for (uint32_t i = 0; i < count; ++i) { quantize(vec + i * dim, dim, ovec + i * out_stride); } @@ -202,7 +228,9 @@ class UniformInt8StreamingReformer : public IndexReformer { return; } for (size_t i = 0; i < dim; ++i) { - float v = std::round(in[i] * scale_ + bias_); + // Match _mm512_cvtps_epi32 used by the SIMD kernel under the default + // floating-point rounding mode. + float v = std::nearbyint(in[i] * scale_ + bias_); v = std::max(0.0f, std::min(127.0f, v)); out[i] = static_cast(v); } @@ -216,8 +244,8 @@ class UniformInt8StreamingReformer : public IndexReformer { turbo::UniformQuantizeFunc quantize_func_{nullptr}; }; -INDEX_FACTORY_REGISTER_REFORMER_ALIAS(UniformInt8StreamingReformer, - UniformInt8StreamingReformer, +INDEX_FACTORY_REGISTER_REFORMER_ALIAS(UniformUint7StreamingReformer, + UniformUint7StreamingReformer, IndexMeta::DataType::DT_INT8); } // namespace core diff --git a/src/core/quantizer/uniform_uint8_converter.cc b/src/core/quantizer/uniform_uint8_converter.cc new file mode 100644 index 000000000..5e37ed5a2 --- /dev/null +++ b/src/core/quantizer/uniform_uint8_converter.cc @@ -0,0 +1,331 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "../metric/metric_params.h" + +namespace zvec { +namespace core { + +class UniformUint8StreamingConverter : public IndexConverter { + public: + explicit UniformUint8StreamingConverter( + IndexMeta::DataType /*destination_type*/) {} + + int init(const IndexMeta &index_meta, const ailego::Params ¶ms) override { + meta_ = index_meta; + original_dimension_ = index_meta.dimension(); + encoded_dimension_ = 0; + scale_ = 0.0f; + bias_ = 0.0f; + holder_.reset(); + stats_.set_trained_count(0); + stats_.set_transformed_count(0); + + if (original_dimension_ == 0 || original_dimension_ > MAX_DIMENSION) { + LOG_ERROR( + "UniformUint8StreamingConverter: dimension=%zu must be in [1, %d]", + original_dimension_, MAX_DIMENSION); + return IndexError_InvalidArgument; + } + encoded_dimension_ = original_dimension_ + kTailBytes; + + meta_.set_converter("UniformUint8StreamingConverter", 0, params); + meta_.set_meta(IndexMeta::DataType::DT_INT8, encoded_dimension_); + + ailego::Params metric_params; + metric_params.set(UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME, + index_meta.metric_name()); + meta_.set_metric("UniformUint8", 0, metric_params); + + const bool has_scale = params.get(UNIFORM_UINT8_REFORMER_SCALE, &scale_); + const bool has_bias = params.get(UNIFORM_UINT8_REFORMER_BIAS, &bias_); + if (has_scale != has_bias || + (has_scale && + (!std::isfinite(scale_) || scale_ <= 0.0f || !std::isfinite(bias_)))) { + LOG_ERROR("UniformUint8StreamingConverter: invalid scale/bias params"); + return IndexError_InvalidArgument; + } + + if (has_scale) { + ailego::Params reformer_params; + reformer_params.set(UNIFORM_UINT8_REFORMER_SCALE, scale_); + reformer_params.set(UNIFORM_UINT8_REFORMER_BIAS, bias_); + meta_.set_reformer("UniformUint8StreamingReformer", 0, reformer_params); + } + return 0; + } + + int cleanup(void) override { + stats_.set_trained_count(0); + stats_.set_transformed_count(0); + holder_.reset(); + scale_ = 0.0f; + bias_ = 0.0f; + return 0; + } + + int train(IndexHolder::Pointer holder) override { + if (!holder) { + LOG_ERROR("UniformUint8StreamingConverter: null training holder"); + return IndexError_InvalidArgument; + } + if (holder->data_type() != IndexMeta::DataType::DT_FP32 || + holder->dimension() != original_dimension_) { + return IndexError_Mismatch; + } + + ailego::ElapsedTime timer; + AILEGO_DEFER([&]() { stats_.set_trained_costtime(timer.milli_seconds()); }); + stats_.set_trained_count(0); + + float global_min = std::numeric_limits::max(); + float global_max = std::numeric_limits::lowest(); + bool all_integer = true; + + auto iterator = holder->create_iterator(); + if (!iterator) { + LOG_ERROR("UniformUint8StreamingConverter: failed to create iterator"); + return IndexError_Runtime; + } + + for (; iterator->is_valid(); iterator->next()) { + const auto *vector = static_cast(iterator->data()); + for (size_t i = 0; i < original_dimension_; ++i) { + const float value = vector[i]; + if (!std::isfinite(value)) { + LOG_ERROR( + "UniformUint8StreamingConverter: non-finite training value " + "(record_idx=%zu, dim_idx=%zu, value=%f)", + static_cast(stats_.trained_count()), i, value); + return IndexError_InvalidArgument; + } + global_min = std::min(global_min, value); + global_max = std::max(global_max, value); + if (all_integer && std::floor(value) != value) { + all_integer = false; + } + } + stats_.set_trained_count(stats_.trained_count() + 1); + } + + if (stats_.trained_count() == 0) { + LOG_ERROR("UniformUint8StreamingConverter: empty training set"); + return IndexError_InvalidArgument; + } + + const float range = global_max - global_min; + if (!std::isfinite(range)) { + LOG_ERROR("UniformUint8StreamingConverter: non-finite training range"); + return IndexError_InvalidArgument; + } + if (all_integer && range <= 255.0f) { + scale_ = 1.0f; + bias_ = -global_min; + } else { + constexpr float epsilon = std::numeric_limits::epsilon(); + scale_ = 255.0f / std::max(range, epsilon); + bias_ = -global_min * scale_; + } + + ailego::Params reformer_params; + reformer_params.set(UNIFORM_UINT8_REFORMER_SCALE, scale_); + reformer_params.set(UNIFORM_UINT8_REFORMER_BIAS, bias_); + meta_.set_reformer("UniformUint8StreamingReformer", 0, reformer_params); + + ailego::Params converter_params = meta_.converter_params(); + converter_params.set(UNIFORM_UINT8_REFORMER_SCALE, scale_); + converter_params.set(UNIFORM_UINT8_REFORMER_BIAS, bias_); + meta_.set_converter(meta_.converter_name(), meta_.converter_revision(), + converter_params); + + LOG_INFO( + "UniformUint8StreamingConverter trained: count=%zu min=%f max=%f " + "scale=%f bias=%f cost=%zums", + static_cast(stats_.trained_count()), global_min, global_max, + scale_, bias_, static_cast(timer.milli_seconds())); + return 0; + } + + int transform(IndexHolder::Pointer holder) override { + if (!holder || holder->data_type() != IndexMeta::DataType::DT_FP32 || + holder->dimension() != original_dimension_ || !std::isfinite(scale_) || + scale_ <= 0.0f || !std::isfinite(bias_)) { + return IndexError_Mismatch; + } + stats_.set_transformed_count(stats_.transformed_count() + holder->count()); + holder_ = std::make_shared( + std::move(holder), original_dimension_, encoded_dimension_, scale_, + bias_); + return 0; + } + + int dump(const IndexDumper::Pointer & /*dumper*/) override { + return 0; + } + + const Stats &stats(void) const override { + return stats_; + } + + IndexHolder::Pointer result(void) const override { + return holder_; + } + + const IndexMeta &meta(void) const override { + return meta_; + } + + private: + // Stored squared-sum tail and exact squared L2. + static_assert(uint64_t{MAX_DIMENSION} * 255 * 255 <= + (std::numeric_limits::max)()); + // Signed dot product accumulated by the VNNI query kernel. + static_assert(uint64_t{MAX_DIMENSION} * 255 * 128 <= + (std::numeric_limits::max)()); + // Absolute value of sum(q*q) - 256*sum(q), maximized at q=128. + static_assert(uint64_t{MAX_DIMENSION} * 128 * 128 <= + (std::numeric_limits::max)()); + static constexpr size_t kTailBytes = sizeof(uint32_t); + + static void EncodeRecord(const float *input, size_t dimension, float scale, + float bias, int8_t *output) { + auto *bytes = reinterpret_cast(output); + for (size_t i = 0; i < dimension; ++i) { + float value = std::round(input[i] * scale + bias); + value = std::max(0.0f, std::min(255.0f, value)); + bytes[i] = static_cast(value); + } + + int64_t sum_squared = 0; + for (size_t i = 0; i < dimension; ++i) { + const int code = static_cast(bytes[i]); + sum_squared += code * code; + bytes[i] = static_cast(code - 128); + } + const uint32_t tail = static_cast(sum_squared); + std::memcpy(output + dimension, &tail, sizeof(tail)); + } + + class UniformUint8Holder : public IndexHolder { + public: + class Iterator : public IndexHolder::Iterator { + public: + Iterator(const UniformUint8Holder *owner, + IndexHolder::Iterator::Pointer iterator) + : owner_(owner), + buffer_(owner->encoded_dimension_, 0), + iterator_(std::move(iterator)) { + encode(); + } + + const void *data(void) const override { + return buffer_.data(); + } + + bool is_valid(void) const override { + return iterator_ && iterator_->is_valid(); + } + + uint64_t key(void) const override { + return iterator_->key(); + } + + void next(void) override { + iterator_->next(); + encode(); + } + + private: + void encode() { + if (!is_valid()) { + return; + } + EncodeRecord(static_cast(iterator_->data()), + owner_->original_dimension_, owner_->scale_, owner_->bias_, + buffer_.data()); + } + + const UniformUint8Holder *owner_; + std::vector buffer_; + IndexHolder::Iterator::Pointer iterator_; + }; + + UniformUint8Holder(IndexHolder::Pointer holder, size_t original_dimension, + size_t encoded_dimension, float scale, float bias) + : holder_(std::move(holder)), + original_dimension_(original_dimension), + encoded_dimension_(encoded_dimension), + scale_(scale), + bias_(bias) {} + + size_t count(void) const override { + return holder_->count(); + } + + size_t dimension(void) const override { + return encoded_dimension_; + } + + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_INT8; + } + + size_t element_size(void) const override { + return IndexMeta::ElementSizeof(IndexMeta::DataType::DT_INT8, + encoded_dimension_); + } + + bool multipass(void) const override { + return holder_->multipass(); + } + + IndexHolder::Iterator::Pointer create_iterator(void) override { + auto iterator = holder_->create_iterator(); + return iterator ? std::make_unique(this, std::move(iterator)) + : nullptr; + } + + private: + IndexHolder::Pointer holder_; + size_t original_dimension_; + size_t encoded_dimension_; + float scale_; + float bias_; + }; + + IndexMeta meta_{}; + Stats stats_{}; + IndexHolder::Pointer holder_{}; + size_t original_dimension_{0}; + size_t encoded_dimension_{0}; + float scale_{0.0f}; + float bias_{0.0f}; +}; + +INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(UniformUint8StreamingConverter, + UniformUint8StreamingConverter, + IndexMeta::DataType::DT_INT8); + +} // namespace core +} // namespace zvec diff --git a/src/core/quantizer/uniform_uint8_reformer.cc b/src/core/quantizer/uniform_uint8_reformer.cc new file mode 100644 index 000000000..71ca180e1 --- /dev/null +++ b/src/core/quantizer/uniform_uint8_reformer.cc @@ -0,0 +1,231 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { + +class UniformUint8StreamingReformer : public IndexReformer { + public: + explicit UniformUint8StreamingReformer( + IndexMeta::DataType /*destination_type*/) {} + + int init(const ailego::Params ¶ms) override { + Reset(); + + float scale = 0.0f; + float bias = 0.0f; + const bool has_scale = params.get(UNIFORM_UINT8_REFORMER_SCALE, &scale); + const bool has_bias = params.get(UNIFORM_UINT8_REFORMER_BIAS, &bias); + if (!has_scale || !has_bias) { + LOG_ERROR( + "UniformUint8StreamingReformer: missing scale/bias params " + "(scale_present=%d, bias_present=%d)", + static_cast(has_scale), static_cast(has_bias)); + return IndexError_InvalidArgument; + } + return SetParams(scale, bias); + } + + int cleanup(void) override { + Reset(); + return 0; + } + + int load(IndexStorage::Pointer) override { + return 0; + } + + int unload(void) override { + return 0; + } + + int transform(const void *query, const IndexQueryMeta &query_meta, + std::string *output, + IndexQueryMeta *output_meta) const override { + return Encode(query, query_meta, 1, output, output_meta, + /*shift_output=*/false); + } + + int transform(const void *queries, const IndexQueryMeta &query_meta, + uint32_t count, std::string *output, + IndexQueryMeta *output_meta) const override { + return Encode(queries, query_meta, count, output, output_meta, + /*shift_output=*/false); + } + + int convert(const void *record, const IndexQueryMeta &record_meta, + std::string *output, IndexQueryMeta *output_meta) const override { + return Encode(record, record_meta, 1, output, output_meta, + /*shift_output=*/true); + } + + int convert(const void *records, const IndexQueryMeta &record_meta, + uint32_t count, std::string *output, + IndexQueryMeta *output_meta) const override { + return Encode(records, record_meta, count, output, output_meta, + /*shift_output=*/true); + } + + int normalize(const void * /*query*/, const IndexQueryMeta & /*query_meta*/, + IndexDocumentList &result) const override { + if (!initialized_) { + return IndexError_Runtime; + } + for (auto &document : result) { + *document.mutable_score() *= scale_reciprocal_squared_; + } + return 0; + } + + bool need_revert(void) const override { + return true; + } + + // Revert only accepts the shifted record layout emitted by convert(). + // transform() emits an unshifted query layout and is not a valid input. + int revert(const void *input, const IndexQueryMeta &record_meta, + std::string *output) const override { + if (!initialized_ || !input || !output || + record_meta.data_type() != IndexMeta::DataType::DT_INT8) { + return IndexError_InvalidArgument; + } + + if (record_meta.dimension() <= kTailBytes) { + return IndexError_InvalidArgument; + } + const size_t dimension = record_meta.dimension() - kTailBytes; + output->resize(dimension * sizeof(float)); + auto *destination = output->data(); + const auto *source = static_cast(input); + const float inverse_scale = 1.0f / scale_; + for (size_t i = 0; i < dimension; ++i) { + const float raw_code = static_cast(source[i]) + 128.0f; + const float value = (raw_code - bias_) * inverse_scale; + std::memcpy(destination + i * sizeof(value), &value, sizeof(value)); + } + return 0; + } + + private: + void Reset() { + initialized_ = false; + scale_ = 0.0f; + bias_ = 0.0f; + scale_reciprocal_squared_ = 1.0f; + } + + int SetParams(float scale, float bias) { + if (!std::isfinite(scale) || scale <= 0.0f || !std::isfinite(bias)) { + LOG_ERROR( + "UniformUint8StreamingReformer: invalid params scale=%f bias=%f", + scale, bias); + initialized_ = false; + return IndexError_InvalidArgument; + } + scale_ = scale; + bias_ = bias; + scale_reciprocal_squared_ = 1.0f / (scale_ * scale_); + initialized_ = true; + return 0; + } + + int Encode(const void *input, const IndexQueryMeta &input_meta, + uint32_t count, std::string *output, IndexQueryMeta *output_meta, + bool shift_output) const { + if (!initialized_ || !input || !output || !output_meta || count == 0) { + return IndexError_InvalidArgument; + } + if (input_meta.data_type() != IndexMeta::DataType::DT_FP32 || + input_meta.unit_size() != + IndexMeta::UnitSizeof(IndexMeta::DataType::DT_FP32)) { + return IndexError_Unsupported; + } + + const size_t dimension = input_meta.dimension(); + if (dimension == 0 || dimension > MAX_DIMENSION) { + LOG_ERROR( + "UniformUint8StreamingReformer: dimension=%zu must be in [1, %d]", + dimension, MAX_DIMENSION); + return IndexError_InvalidArgument; + } + const size_t encoded_dimension = dimension + kTailBytes; + *output_meta = input_meta; + output_meta->set_meta(IndexMeta::DataType::DT_INT8, encoded_dimension); + const size_t output_stride = output_meta->element_size(); + output->resize(static_cast(count) * output_stride); + + const auto *source = static_cast(input); + auto *destination = reinterpret_cast(output->data()); + for (uint32_t i = 0; i < count; ++i) { + EncodeOne(source + static_cast(i) * dimension, dimension, + destination + static_cast(i) * output_stride, + shift_output); + } + return 0; + } + + void EncodeOne(const float *input, size_t dimension, int8_t *output, + bool shift_output) const { + auto *bytes = reinterpret_cast(output); + for (size_t i = 0; i < dimension; ++i) { + float value = std::round(input[i] * scale_ + bias_); + value = std::max(0.0f, std::min(255.0f, value)); + bytes[i] = static_cast(value); + } + + int64_t sum_squared = 0; + for (size_t i = 0; i < dimension; ++i) { + const int code = static_cast(bytes[i]); + sum_squared += code * code; + if (shift_output) { + bytes[i] = static_cast(code - 128); + } + } + const uint32_t tail = static_cast(sum_squared); + std::memcpy(output + dimension, &tail, sizeof(tail)); + } + + // Keep the encoded-tail and exact-query assumptions local to this + // independent encoding entry point. + static_assert(uint64_t{MAX_DIMENSION} * 255 * 255 <= + (std::numeric_limits::max)(), + "Uniform UINT8 norm and exact L2 must fit uint32_t"); + static_assert(uint64_t{MAX_DIMENSION} * 255 * 128 <= + (std::numeric_limits::max)(), + "Uniform UINT8 VNNI dot product must fit int32_t"); + static_assert(uint64_t{MAX_DIMENSION} * 128 * 128 <= + (std::numeric_limits::max)(), + "Uniform UINT8 query correction must fit int32_t"); + static constexpr size_t kTailBytes = sizeof(uint32_t); + float scale_{0.0f}; + float bias_{0.0f}; + float scale_reciprocal_squared_{1.0f}; + bool initialized_{false}; +}; + +INDEX_FACTORY_REGISTER_REFORMER_ALIAS(UniformUint8StreamingReformer, + UniformUint8StreamingReformer, + IndexMeta::DataType::DT_INT8); + +} // namespace core +} // namespace zvec diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 877e6e3a0..88ef23418 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -94,7 +94,11 @@ enum class QuantizerType { kInt8, kInt4, kRabitq, - kUniformInt8, // Global uniform int8 quantization (shared scale/bias). + // Global uniform quantization with a shared scale/bias and 128 code values + // in [0, 127]. + kUniformUint7 = 8, + // Global uniform quantization with the full uint8 code range [0, 255]. + kUniformUint8 = 9, }; struct ZVEC_CORE_API SerializableBase { diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 5baa8c418..c92bc87da 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -50,11 +50,9 @@ using BatchDistanceFunc = std::function int8 with a global affine transform: -// out[i] = clip(round(in[i] * scale + bias), 0, 127) -// This signature is specific to the uniform-int8 quantizer and is NOT a -// generic quantize contract. Raw function pointer (rather than std::function) -// to avoid indirect-call overhead on the per-record / per-query hot path. +// Uniform UINT7 quantize kernel: fp32 -> int8 code in [0, 127] with a global +// affine transform. Raw function pointer (rather than std::function) avoids +// indirect-call overhead on the per-record / per-query hot path. using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale, float bias, int8_t *out); @@ -89,12 +87,13 @@ enum class DataType { enum class QuantizeType { kDefault, - kUniform, + kUniform, // Uniform uint7: codes are restricted to [0, 127]. kRecord, kFp16, kFp32, kPQ, - kRabit + kRabit, + kUniformUint8, // Uniform uint8: codes cover the full [0, 255] range. }; enum class RotateType : uint16_t { diff --git a/src/turbo/distance/avx512_vnni/uniform_int8/quantize.cc b/src/turbo/distance/avx512_vnni/uniform_uint7/quantize.cc similarity index 81% rename from src/turbo/distance/avx512_vnni/uniform_int8/quantize.cc rename to src/turbo/distance/avx512_vnni/uniform_uint7/quantize.cc index 140923a23..ae52018a5 100644 --- a/src/turbo/distance/avx512_vnni/uniform_int8/quantize.cc +++ b/src/turbo/distance/avx512_vnni/uniform_uint7/quantize.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -// AVX-512 quantization for the uniform-int8 quantizer. +// AVX-512 quantization for the uniform-uint7 quantizer. // // Pipeline (16 floats per iteration): // 1. Load 16 fp32 values (vmovups) @@ -27,7 +27,7 @@ // // Compiled with -march=avx512vnni (set per-file in src/turbo/CMakeLists.txt). -#include "avx512_vnni/uniform_int8/quantize.h" +#include "avx512_vnni/uniform_uint7/quantize.h" #include #include @@ -36,8 +36,8 @@ namespace zvec::turbo::avx512_vnni { -void uniform_int8_quantize(const float *in, std::size_t dim, float scale, - float bias, std::int8_t *out) { +void uniform_uint7_quantize(const float *in, std::size_t dim, float scale, + float bias, std::int8_t *out) { const __m512 vscale = _mm512_set1_ps(scale); const __m512 vbias = _mm512_set1_ps(bias); const __m512i vzero = _mm512_setzero_si512(); @@ -57,9 +57,10 @@ void uniform_int8_quantize(const float *in, std::size_t dim, float scale, _mm_storeu_si128(reinterpret_cast<__m128i *>(out + i), packed); } - // Tail: scalar fallback (matches the scalar reference exactly). + // Tail: use the same rounding mode as _mm512_cvtps_epi32 so dimensions that + // are not multiples of 16 remain bit-exact with the vectorized prefix. for (; i < dim; ++i) { - float v = std::round(in[i] * scale + bias); + float v = std::nearbyint(in[i] * scale + bias); v = std::max(0.0f, std::min(127.0f, v)); out[i] = static_cast(v); } @@ -71,9 +72,9 @@ void uniform_int8_quantize(const float *in, std::size_t dim, float scale, namespace zvec::turbo::avx512_vnni { -void uniform_int8_quantize(const float * /*in*/, std::size_t /*dim*/, - float /*scale*/, float /*bias*/, - std::int8_t * /*out*/) { +void uniform_uint7_quantize(const float * /*in*/, std::size_t /*dim*/, + float /*scale*/, float /*bias*/, + std::int8_t * /*out*/) { // Intentionally empty; turbo::get_uniform_quantize_func will return nullptr // on CPUs without AVX-512 support and the caller will use its scalar path. } diff --git a/src/turbo/distance/avx512_vnni/uniform_int8/quantize.h b/src/turbo/distance/avx512_vnni/uniform_uint7/quantize.h similarity index 63% rename from src/turbo/distance/avx512_vnni/uniform_int8/quantize.h rename to src/turbo/distance/avx512_vnni/uniform_uint7/quantize.h index f544711a1..d7133287e 100644 --- a/src/turbo/distance/avx512_vnni/uniform_int8/quantize.h +++ b/src/turbo/distance/avx512_vnni/uniform_uint7/quantize.h @@ -19,15 +19,15 @@ namespace zvec::turbo::avx512_vnni { -// AVX-512 vectorized quantization for the uniform-int8 quantizer. -// forward: out[i] = clip(round(in[i] * scale + bias), -127, 127) +// AVX-512 vectorized quantization for the uniform-uint7 quantizer. +// forward: +// out[i] = clip(round-to-nearest-even(in[i] * scale + bias), 0, 127) // // Implementation detail: relies on hardware saturation in // vcvtsepi32_epi8 / vpackss to clip without explicit min/max. -// Note: AVX-512 default rounding mode is round-to-nearest-even, which -// matches std::round() to within ULP for typical embedding values; tests -// against the scalar reference confirm bit-exact results on common inputs. -void uniform_int8_quantize(const float *in, std::size_t dim, float scale, - float bias, std::int8_t *out); +// Both the vectorized loop and its scalar tail honor the active floating-point +// rounding mode; under the default mode, halfway values round to nearest even. +void uniform_uint7_quantize(const float *in, std::size_t dim, float scale, + float bias, std::int8_t *out); } // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/uniform_uint7/squared_euclidean.cc similarity index 87% rename from src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.cc rename to src/turbo/distance/avx512_vnni/uniform_uint7/squared_euclidean.cc index 1d6c0a0f4..b100d7354 100644 --- a/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.cc +++ b/src/turbo/distance/avx512_vnni/uniform_uint7/squared_euclidean.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -// AVX512-VNNI optimized squared Euclidean distance for uniform-quantized INT8. +// AVX512-VNNI optimized squared Euclidean distance for uniform uint7 codes. // // Since all vectors share a single global scale/bias, the distance is simply: // sum((a[i] - b[i])^2) @@ -35,7 +35,7 @@ // This file is compiled with per-file -march=avx512vnni (set in // CMakeLists.txt). -#include "avx512_vnni/uniform_int8/squared_euclidean.h" +#include "avx512_vnni/uniform_uint7/squared_euclidean.h" #include "zvec/ailego/internal/platform.h" #if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) @@ -115,10 +115,10 @@ static ailego_force_inline void uniform_sq_l2_int8_batch_impl( } // --------------------------------------------------------------------------- -// Public: single-vector squared Euclidean distance (int8, VNNI abs trick) +// Public: single-vector squared Euclidean distance (uint7, VNNI abs trick) // --------------------------------------------------------------------------- -void uniform_squared_euclidean_int8_distance(const void *a, const void *b, - size_t dim, float *distance) { +void uniform_squared_euclidean_uint7_distance(const void *a, const void *b, + size_t dim, float *distance) { const int8_t *lhs = reinterpret_cast(a); const int8_t *rhs = reinterpret_cast(b); @@ -174,12 +174,12 @@ void uniform_squared_euclidean_int8_distance(const void *a, const void *b, } // --------------------------------------------------------------------------- -// Public: batch squared Euclidean distance (int8, no tail, no preprocessing) +// Public: batch squared Euclidean distance (uint7, no tail, no preprocessing) // --------------------------------------------------------------------------- -void uniform_squared_euclidean_int8_batch_distance(const void *const *vectors, - const void *query, size_t n, - size_t dim, - float *distances) { +void uniform_squared_euclidean_uint7_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, + float *distances) { static constexpr size_t batch_size = 4; static constexpr size_t prefetch_step = 2; @@ -198,8 +198,8 @@ void uniform_squared_euclidean_int8_batch_distance(const void *const *vectors, // both an extra `batch_size=1` template instantiation and the per-call // std::array setup that the batch_impl path requires. for (; i < n; ++i) { - uniform_squared_euclidean_int8_distance(vectors[i], query, dim, - distances + i); + uniform_squared_euclidean_uint7_distance(vectors[i], query, dim, + distances + i); } } @@ -209,11 +209,12 @@ void uniform_squared_euclidean_int8_batch_distance(const void *const *vectors, namespace zvec::turbo::avx512_vnni { -void uniform_squared_euclidean_int8_distance(const void * /*a*/, - const void * /*b*/, size_t /*dim*/, - float * /*distance*/) {} +void uniform_squared_euclidean_uint7_distance(const void * /*a*/, + const void * /*b*/, + size_t /*dim*/, + float * /*distance*/) {} -void uniform_squared_euclidean_int8_batch_distance( +void uniform_squared_euclidean_uint7_batch_distance( const void *const * /*vectors*/, const void * /*query*/, size_t /*n*/, size_t /*dim*/, float * /*distances*/) {} diff --git a/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.h b/src/turbo/distance/avx512_vnni/uniform_uint7/squared_euclidean.h similarity index 58% rename from src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.h rename to src/turbo/distance/avx512_vnni/uniform_uint7/squared_euclidean.h index 14bad3f2d..707a8483d 100644 --- a/src/turbo/distance/avx512_vnni/uniform_int8/squared_euclidean.h +++ b/src/turbo/distance/avx512_vnni/uniform_uint7/squared_euclidean.h @@ -18,18 +18,18 @@ namespace zvec::turbo::avx512_vnni { -// Compute squared Euclidean distance between two uniform-quantized INT8 -// vectors. Unlike record_quantized, there is NO metadata tail — `dim` is the -// pure int8 vector length. Distance = sum((a[i] - b[i])^2). -void uniform_squared_euclidean_int8_distance(const void *a, const void *b, - size_t dim, float *distance); +// Compute squared Euclidean distance between two uniform-quantized UINT7 +// vectors. Unlike record_quantized, there is no metadata tail: `dim` is the +// pure int8 vector length. Distance = sum((a[i] - b[i])^2). +void uniform_squared_euclidean_uint7_distance(const void *a, const void *b, + size_t dim, float *distance); -// Batch version: compute squared Euclidean distance between `n` INT8 database -// vectors and a single INT8 query. No query preprocessing is required (unlike +// Batch version: compute squared Euclidean distance between `n` UINT7 database +// vectors and a single UINT7 query. No query preprocessing is required (unlike // the record_quantized path which needs int8→uint8 shifting for dpbusd). -void uniform_squared_euclidean_int8_batch_distance(const void *const *vectors, - const void *query, size_t n, - size_t dim, - float *distances); +void uniform_squared_euclidean_uint7_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, + float *distances); } // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.cc new file mode 100644 index 000000000..5244d8d17 --- /dev/null +++ b/src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.cc @@ -0,0 +1,462 @@ +// 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. + +// AVX512-VNNI optimized squared L2 for uniform uint8 quantization. +// +// Stored record layout: [dim int8(code - 128) | uint32 sum_sq(raw code)]. +// Query layout: [dim raw uint8 code | int32 query correction]. +// +// Build-time pairwise distance uses true L2 between two stored shifted +// vectors. Search-time batch distance includes the query-only correction: +// distance = sum_sq(record_raw) +// - 2 * dot(record_shifted, query_raw) +// + sum_sq(query_raw) - 256 * sum(query_raw) +// = ||record_raw - query_raw||^2 +// VNNI uses vpdpbusd's unsigned x signed contract as: +// dot(record_shifted, query_raw) = dpbusd(query_raw, record_shifted) +// +// Batch kernel design (hot path for graph search): +// - four records per block with independent accumulators +// - software prefetch of future records, including the metadata tail +// - SIMD horizontal reduction and final score calculation for four records + +#include "avx512_vnni/uniform_uint8/squared_euclidean.h" +#include +#include +#include +#include "zvec/ailego/internal/platform.h" + +#if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) +#include +#endif + +namespace zvec::turbo::avx512_vnni { + +namespace { + +constexpr size_t kTailBytes = sizeof(uint32_t); +// The uint32 norm and int32 VNNI dot product are both lossless through the +// public 65,536-dimension limit. Larger direct calls use squared differences. +constexpr size_t kMaxIdentityDimension = 65536; +static_assert(uint64_t{kMaxIdentityDimension} * 255 * 255 <= + (std::numeric_limits::max)()); +static_assert(uint64_t{kMaxIdentityDimension} * 255 * 128 <= + (std::numeric_limits::max)()); +static_assert(uint64_t{kMaxIdentityDimension} * 128 * 128 <= + (std::numeric_limits::max)()); + +static inline size_t original_dim(size_t encoded_dim) { + return encoded_dim > kTailBytes ? encoded_dim - kTailBytes : 0; +} + +static inline void uniform_sq_l2_uint8_scalar_single(const void *vector, + const uint8_t *raw_query, + size_t orig_dim, + float *distance) { + const auto *record = reinterpret_cast(vector); + int64_t result = 0; + for (size_t d = 0; d < orig_dim; ++d) { + const int difference = + static_cast(record[d]) - (static_cast(raw_query[d]) - 128); + result += static_cast(difference) * difference; + } + *distance = static_cast(result); +} + +#if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) + +static inline uint32_t tail(const void *vector, size_t orig_dim) { + uint32_t value = 0; + std::memcpy(&value, reinterpret_cast(vector) + orig_dim, + sizeof(value)); + return value; +} + +static inline int32_t query_correction(const void *query, size_t orig_dim) { + int32_t value = 0; + std::memcpy(&value, reinterpret_cast(query) + orig_dim, + sizeof(value)); + return value; +} + +// Convert the low four uint32 lanes with AVX512F's native unsigned +// conversion. The upper lanes are unused. +static ailego_force_inline __m128 uint32_to_float(__m128i values) { + const __m512 converted = _mm512_cvtepu32_ps(_mm512_castsi128_si512(values)); + return _mm512_castps512_ps128(converted); +} + +// Sign-extend 32 stored int8 values, subtract without int8 overflow, and use +// VNNI's signed-word dot product to accumulate pairs of squared differences +// into 16 int32 lanes. The full stored range [-128, 127] produces differences +// in [-255, 255], so each int16 product remains exact. +static ailego_force_inline __m512i squared_diff_32(__m512i accumulator, + const int8_t *lhs, + const int8_t *rhs) { + const __m512i lhs16 = _mm512_cvtepi8_epi16( + _mm256_loadu_si256(reinterpret_cast(lhs))); + const __m512i rhs16 = _mm512_cvtepi8_epi16( + _mm256_loadu_si256(reinterpret_cast(rhs))); + const __m512i difference16 = _mm512_sub_epi16(lhs16, rhs16); + return _mm512_dpwssd_epi32(accumulator, difference16, difference16); +} + +static ailego_force_inline __m512i squared_diff_masked_32(__m512i accumulator, + const int8_t *lhs, + const int8_t *rhs, + __mmask32 mask) { + const __m512i lhs16 = _mm512_cvtepi8_epi16( + _mm256_maskz_loadu_epi8(mask, static_cast(lhs))); + const __m512i rhs16 = _mm512_cvtepi8_epi16( + _mm256_maskz_loadu_epi8(mask, static_cast(rhs))); + const __m512i difference16 = _mm512_sub_epi16(lhs16, rhs16); + return _mm512_dpwssd_epi32(accumulator, difference16, difference16); +} + +// Widen before the horizontal sum. A uint8 squared distance can exceed +// INT32_MAX (for example 65,536 * 255^2), even though every accumulator lane +// remains in range between periodic flushes. +static ailego_force_inline int64_t +reduce_add_epi32_to_int64(__m512i accumulator) { + const __m256i low32 = _mm512_castsi512_si256(accumulator); + const __m256i high32 = _mm512_extracti64x4_epi64(accumulator, 1); + const __m512i low64 = _mm512_cvtepi32_epi64(low32); + const __m512i high64 = _mm512_cvtepi32_epi64(high32); + return _mm512_reduce_add_epi64(_mm512_add_epi64(low64, high64)); +} + +#endif + +} // namespace + +void uniform_squared_euclidean_uint8_distance(const void *a, const void *b, + size_t dim, float *distance) { + const size_t orig_dim = original_dim(dim); + if (orig_dim == 0) { + *distance = 0.0f; + return; + } + + const auto *lhs = reinterpret_cast(a); + const auto *rhs = reinterpret_cast(b); + +#if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) + // Four dependency chains cover 128 bytes per iteration. Each VPDPWSSD lane + // receives two squares, whose maximum contribution is 2 * 255^2. Flush + // every 8,192 iterations so each int32 lane stays below 1.1 billion. + constexpr size_t kBlockBytes = 32; + constexpr size_t kUnrolledBytes = 4 * kBlockBytes; + constexpr size_t kFlushIterations = 8192; + + __m512i accumulator0 = _mm512_setzero_si512(); + __m512i accumulator1 = _mm512_setzero_si512(); + __m512i accumulator2 = _mm512_setzero_si512(); + __m512i accumulator3 = _mm512_setzero_si512(); + int64_t result = 0; + + size_t d = 0; + size_t iterations_since_flush = 0; + for (; d + kUnrolledBytes <= orig_dim; d += kUnrolledBytes) { + accumulator0 = squared_diff_32(accumulator0, lhs + d, rhs + d); + accumulator1 = squared_diff_32(accumulator1, lhs + d + 32, rhs + d + 32); + accumulator2 = squared_diff_32(accumulator2, lhs + d + 64, rhs + d + 64); + accumulator3 = squared_diff_32(accumulator3, lhs + d + 96, rhs + d + 96); + + if (++iterations_since_flush == kFlushIterations) { + result += reduce_add_epi32_to_int64(accumulator0); + result += reduce_add_epi32_to_int64(accumulator1); + result += reduce_add_epi32_to_int64(accumulator2); + result += reduce_add_epi32_to_int64(accumulator3); + accumulator0 = _mm512_setzero_si512(); + accumulator1 = _mm512_setzero_si512(); + accumulator2 = _mm512_setzero_si512(); + accumulator3 = _mm512_setzero_si512(); + iterations_since_flush = 0; + } + } + + for (; d + kBlockBytes <= orig_dim; d += kBlockBytes) { + accumulator0 = squared_diff_32(accumulator0, lhs + d, rhs + d); + } + + if (d < orig_dim) { + const size_t remaining = orig_dim - d; + const __mmask32 mask = + static_cast<__mmask32>((uint32_t{1} << remaining) - 1); + accumulator0 = squared_diff_masked_32(accumulator0, lhs + d, rhs + d, mask); + } + + result += reduce_add_epi32_to_int64(accumulator0); + result += reduce_add_epi32_to_int64(accumulator1); + result += reduce_add_epi32_to_int64(accumulator2); + result += reduce_add_epi32_to_int64(accumulator3); + *distance = static_cast(result); +#else + int64_t result = 0; + for (size_t i = 0; i < orig_dim; ++i) { + const int difference = static_cast(lhs[i]) - static_cast(rhs[i]); + result += static_cast(difference) * difference; + } + *distance = static_cast(result); +#endif +} + +#if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) + +namespace { + +// Reduce four zmm int32 accumulators to one xmm containing four scalar sums. +static ailego_force_inline __m128i reduce_add_4x16_epi32(__m512i accumulator0, + __m512i accumulator1, + __m512i accumulator2, + __m512i accumulator3) { + const __m256i half0 = + _mm256_add_epi32(_mm512_castsi512_si256(accumulator0), + _mm512_extracti64x4_epi64(accumulator0, 1)); + const __m256i half1 = + _mm256_add_epi32(_mm512_castsi512_si256(accumulator1), + _mm512_extracti64x4_epi64(accumulator1, 1)); + const __m256i half2 = + _mm256_add_epi32(_mm512_castsi512_si256(accumulator2), + _mm512_extracti64x4_epi64(accumulator2, 1)); + const __m256i half3 = + _mm256_add_epi32(_mm512_castsi512_si256(accumulator3), + _mm512_extracti64x4_epi64(accumulator3, 1)); + const __m256i pair01 = _mm256_hadd_epi32(half0, half1); + const __m256i pair23 = _mm256_hadd_epi32(half2, half3); + const __m256i totals = _mm256_hadd_epi32(pair01, pair23); + return _mm_add_epi32(_mm256_castsi256_si128(totals), + _mm256_extracti128_si256(totals, 1)); +} + +static ailego_force_inline void uniform_sq_l2_uint8_batch4( + const void *const *vectors, const uint8_t *raw_query, size_t orig_dim, + int32_t correction, const void *const *prefetch_vectors, float *distances) { + __m512i accumulator0 = _mm512_setzero_si512(); + __m512i accumulator1 = _mm512_setzero_si512(); + __m512i accumulator2 = _mm512_setzero_si512(); + __m512i accumulator3 = _mm512_setzero_si512(); + + const auto *vector0 = reinterpret_cast(vectors[0]); + const auto *vector1 = reinterpret_cast(vectors[1]); + const auto *vector2 = reinterpret_cast(vectors[2]); + const auto *vector3 = reinterpret_cast(vectors[3]); + + size_t d = 0; + for (; d + 64 <= orig_dim; d += 64) { + const __m512i query = + _mm512_loadu_si512(reinterpret_cast(raw_query + d)); + const __m512i record0 = + _mm512_loadu_si512(reinterpret_cast(vector0 + d)); + const __m512i record1 = + _mm512_loadu_si512(reinterpret_cast(vector1 + d)); + const __m512i record2 = + _mm512_loadu_si512(reinterpret_cast(vector2 + d)); + const __m512i record3 = + _mm512_loadu_si512(reinterpret_cast(vector3 + d)); + + for (size_t i = 0; i < 4; ++i) { + if (prefetch_vectors[i]) { + _mm_prefetch(reinterpret_cast(prefetch_vectors[i]) + d, + _MM_HINT_T0); + } + } + + accumulator0 = _mm512_dpbusd_epi32(accumulator0, query, record0); + accumulator1 = _mm512_dpbusd_epi32(accumulator1, query, record1); + accumulator2 = _mm512_dpbusd_epi32(accumulator2, query, record2); + accumulator3 = _mm512_dpbusd_epi32(accumulator3, query, record3); + } + + // The main loop only covers full cache lines, so prefetch the metadata tail + // of each future record explicitly. + for (size_t i = 0; i < 4; ++i) { + if (prefetch_vectors[i]) { + _mm_prefetch( + reinterpret_cast(prefetch_vectors[i]) + orig_dim, + _MM_HINT_T0); + } + } + + __m128i dot_products = reduce_add_4x16_epi32(accumulator0, accumulator1, + accumulator2, accumulator3); + + if (d < orig_dim) { + alignas(16) int32_t totals[4]; + _mm_store_si128(reinterpret_cast<__m128i *>(totals), dot_products); + const int8_t *records[4] = {vector0, vector1, vector2, vector3}; + for (size_t i = 0; i < 4; ++i) { + int32_t remainder = 0; + for (size_t j = d; j < orig_dim; ++j) { + remainder += + static_cast(records[i][j]) * static_cast(raw_query[j]); + } + totals[i] += remainder; + } + dot_products = _mm_load_si128(reinterpret_cast(totals)); + } + + alignas(16) const uint32_t tails[4] = { + tail(vector0, orig_dim), tail(vector1, orig_dim), tail(vector2, orig_dim), + tail(vector3, orig_dim)}; + const __m128i sum_squared = + _mm_load_si128(reinterpret_cast(tails)); + + // Exact squared L2 can exceed INT32_MAX while still fitting uint32_t for + // every supported dimension. Preserve its exact bit pattern in packed + // arithmetic, then convert it as unsigned. + const __m128i squared_distances = + _mm_add_epi32(_mm_sub_epi32(sum_squared, _mm_slli_epi32(dot_products, 1)), + _mm_set1_epi32(correction)); + _mm_storeu_ps(distances, uint32_to_float(squared_distances)); +} + +static ailego_force_inline void uniform_sq_l2_uint8_single( + const void *vector, const uint8_t *raw_query, size_t orig_dim, + int32_t correction, float *distance) { + const auto *record = reinterpret_cast(vector); + __m512i accumulator = _mm512_setzero_si512(); + size_t d = 0; + for (; d + 64 <= orig_dim; d += 64) { + const __m512i query = + _mm512_loadu_si512(reinterpret_cast(raw_query + d)); + const __m512i stored = + _mm512_loadu_si512(reinterpret_cast(record + d)); + accumulator = _mm512_dpbusd_epi32(accumulator, query, stored); + } + int64_t dot_product = _mm512_reduce_add_epi32(accumulator); + for (; d < orig_dim; ++d) { + dot_product += static_cast(record[d]) * static_cast(raw_query[d]); + } + *distance = static_cast(static_cast(tail(vector, orig_dim)) - + 2 * dot_product + correction); +} + +} // namespace + +#endif + +void uniform_squared_euclidean_uint8_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, + float *distances) { + const size_t orig_dim = original_dim(dim); + if (orig_dim == 0) { + for (size_t i = 0; i < n; ++i) { + distances[i] = 0.0f; + } + return; + } + if (orig_dim > kMaxIdentityDimension) { + const auto *raw_query = reinterpret_cast(query); + for (size_t i = 0; i < n; ++i) { + uniform_sq_l2_uint8_scalar_single(vectors[i], raw_query, orig_dim, + distances + i); + } + return; + } + +#if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) + const auto *raw_query = reinterpret_cast(query); + const int32_t correction = query_correction(query, orig_dim); + + constexpr size_t kBatchSize = 4; + const size_t prefetch_step = orig_dim > 256 ? 1 : 2; + size_t i = 0; + const void *prefetch_vectors[kBatchSize]; + for (; i + kBatchSize <= n; i += kBatchSize) { + for (size_t j = 0; j < kBatchSize; ++j) { + const size_t prefetch_index = i + j + kBatchSize * prefetch_step; + prefetch_vectors[j] = + prefetch_index < n ? vectors[prefetch_index] : nullptr; + } + uniform_sq_l2_uint8_batch4(vectors + i, raw_query, orig_dim, correction, + prefetch_vectors, distances + i); + } + for (; i < n; ++i) { + uniform_sq_l2_uint8_single(vectors[i], raw_query, orig_dim, correction, + distances + i); + } +#else + const auto *raw_query = reinterpret_cast(query); + for (size_t i = 0; i < n; ++i) { + uniform_sq_l2_uint8_scalar_single(vectors[i], raw_query, orig_dim, + distances + i); + } +#endif +} + +void uniform_squared_euclidean_uint8_query_preprocess(void *query, size_t dim) { + const size_t orig_dim = original_dim(dim); + if (orig_dim == 0) { + return; + } + + const auto *raw_query = reinterpret_cast(query); + // The tail may already contain a correction from an earlier call. Derive + // everything from the immutable query body so preprocessing is idempotent. + uint64_t sum = 0; + uint64_t sum_squared = 0; + size_t d = 0; + +#if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) + const __m512i zero = _mm512_setzero_si512(); + __m512i sums = _mm512_setzero_si512(); + __m512i squared_sums = _mm512_setzero_si512(); + size_t iterations_since_flush = 0; + constexpr size_t kSquaredSumFlushIterations = 4096; + for (; d + 64 <= orig_dim; d += 64) { + const __m512i values = _mm512_loadu_si512(raw_query + d); + sums = _mm512_add_epi64(sums, _mm512_sad_epu8(values, zero)); + const __m512i low_values = + _mm512_cvtepu8_epi16(_mm512_castsi512_si256(values)); + const __m512i high_values = + _mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(values, 1)); + squared_sums = _mm512_dpwssd_epi32(squared_sums, low_values, low_values); + squared_sums = _mm512_dpwssd_epi32(squared_sums, high_values, high_values); + if (++iterations_since_flush == kSquaredSumFlushIterations) { + sum_squared += + static_cast(reduce_add_epi32_to_int64(squared_sums)); + squared_sums = _mm512_setzero_si512(); + iterations_since_flush = 0; + } + } + alignas(64) uint64_t lanes[8]; + _mm512_store_si512(reinterpret_cast<__m512i *>(lanes), sums); + for (uint64_t lane : lanes) { + sum += lane; + } + sum_squared += static_cast(reduce_add_epi32_to_int64(squared_sums)); +#endif + + for (; d < orig_dim; ++d) { + const uint64_t value = raw_query[d]; + sum += value; + sum_squared += value * value; + } + + const int64_t correction = + static_cast(sum_squared) - 256 * static_cast(sum); + if (correction < (std::numeric_limits::min)() || + correction > (std::numeric_limits::max)()) { + // The public quantizer dimension bound keeps the correction in int32. + // Oversized direct calls use the scalar squared-difference fallback. + return; + } + const int32_t encoded_correction = static_cast(correction); + std::memcpy(static_cast(query) + orig_dim, &encoded_correction, + sizeof(encoded_correction)); +} + +} // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.h b/src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.h new file mode 100644 index 000000000..d47c18722 --- /dev/null +++ b/src/turbo/distance/avx512_vnni/uniform_uint8/squared_euclidean.h @@ -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. + +#pragma once + +#include + +namespace zvec::turbo::avx512_vnni { + +// Record layout: +// [ original_dim bytes: int8 values stored as uint8(code) - 128 ] +// [ uint32 sum_sq_u8 ] +// +// Query layout: +// [ original_dim raw uint8 bytes ] +// [ int32 query correction ] +// +// The index data type remains DT_INT8. Build distance computes exact L2 +// between two shifted records. Batch search compares shifted records with a +// preprocessed raw query and returns exact squared L2: +// sum_sq(record_raw) - 2 * dot(record_shifted, query_raw) +// + sum_sq(query_raw) - 256 * sum(query_raw) +void uniform_squared_euclidean_uint8_distance(const void *a, const void *b, + size_t dim, float *distance); + +void uniform_squared_euclidean_uint8_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, + float *distances); + +// Replace the raw query's uint32 squared-sum tail with: +// sum_sq(query_raw) - 256 * sum(query_raw) +// The query body is already raw uint8 and remains unchanged. +void uniform_squared_euclidean_uint8_query_preprocess(void *query, size_t dim); + +} // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 5a6f33d49..511a60586 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -19,8 +19,9 @@ #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 "avx512_vnni/uniform_uint7/quantize.h" +#include "avx512_vnni/uniform_uint7/squared_euclidean.h" +#include "avx512_vnni/uniform_uint8/squared_euclidean.h" #include "neon/rotate/fht/fht.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" @@ -69,7 +70,15 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { - return avx512_vnni::uniform_squared_euclidean_int8_distance; + return avx512_vnni::uniform_squared_euclidean_uint7_distance; + } + } + } + if (quantize_type == QuantizeType::kUniformUint8) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { + if (metric_type == MetricType::kSquaredEuclidean) { + return avx512_vnni::uniform_squared_euclidean_uint8_distance; } } } @@ -112,7 +121,15 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { - return avx512_vnni::uniform_squared_euclidean_int8_batch_distance; + return avx512_vnni::uniform_squared_euclidean_uint7_batch_distance; + } + } + } + if (quantize_type == QuantizeType::kUniformUint8) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { + if (metric_type == MetricType::kSquaredEuclidean) { + return avx512_vnni::uniform_squared_euclidean_uint8_batch_distance; } } } @@ -137,6 +154,13 @@ QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, } } } + if (quantize_type == QuantizeType::kUniformUint8) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI) && + metric_type == MetricType::kSquaredEuclidean) { + return avx512_vnni::uniform_squared_euclidean_uint8_query_preprocess; + } + } } return nullptr; } @@ -147,7 +171,7 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { // AVX512_VNNI flag for now since the kernel lives in the avx512_vnni // directory and is compiled with the same march flag. if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { - return avx512_vnni::uniform_int8_quantize; + return avx512_vnni::uniform_uint7_quantize; } } return nullptr; diff --git a/tests/core/algorithm/vamana/vamana_streamer_test.cc b/tests/core/algorithm/vamana/vamana_streamer_test.cc index 12f833c3f..02cb4670b 100644 --- a/tests/core/algorithm/vamana/vamana_streamer_test.cc +++ b/tests/core/algorithm/vamana/vamana_streamer_test.cc @@ -18,11 +18,15 @@ #include #include #endif +#include +#include #include #include #include +#include #include #include +#include "metric/metric_params.h" #include "tests/test_util.h" #if defined(__GNUC__) || defined(__GNUG__) @@ -39,6 +43,32 @@ namespace core { constexpr size_t kDim = 16; +std::vector EncodeUniformUint8Record( + const std::vector &codes) { + std::vector encoded(codes.size() + sizeof(uint32_t), 0); + uint64_t sum_squared = 0; + for (size_t i = 0; i < codes.size(); ++i) { + encoded[i] = static_cast(static_cast(codes[i]) - 128); + sum_squared += static_cast(codes[i]) * codes[i]; + } + const uint32_t tail = static_cast(sum_squared); + std::memcpy(encoded.data() + codes.size(), &tail, sizeof(tail)); + return encoded; +} + +std::vector EncodeUniformUint8Query(const std::vector &codes) { + std::vector encoded(codes.size() + sizeof(uint32_t), 0); + uint64_t sum_squared = 0; + auto *bytes = reinterpret_cast(encoded.data()); + for (size_t i = 0; i < codes.size(); ++i) { + bytes[i] = codes[i]; + sum_squared += static_cast(codes[i]) * codes[i]; + } + const uint32_t tail = static_cast(sum_squared); + std::memcpy(encoded.data() + codes.size(), &tail, sizeof(tail)); + return encoded; +} + class VamanaStreamerTest : public testing::Test { protected: void SetUp(void) override; @@ -861,6 +891,83 @@ TEST_F(VamanaStreamerTest, TestAsymmetricQueryMetric) { EXPECT_FLOAT_EQ(-2.0f, primary_key_context->result()[0].score()); } +TEST_F(VamanaStreamerTest, TestUniformUint8UsesQueryMetric) { + constexpr size_t kOriginalDimension = 4; + constexpr size_t kEncodedDimension = kOriginalDimension + sizeof(uint32_t); + + ailego::Params metric_params; + metric_params.set(UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME, + std::string("SquaredEuclidean")); + IndexMeta meta(IndexMeta::DataType::DT_INT8, kEncodedDimension); + meta.set_metric("UniformUint8", 0, metric_params); + + ailego::Params params; + params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 8U); + params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 16U); + params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f); + params.set(PARAM_VAMANA_STREAMER_EF, 16U); + params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 0U); + + auto streamer = IndexFactory::CreateStreamer("VamanaStreamer"); + ASSERT_TRUE(streamer); + ASSERT_EQ(0, streamer->init(meta, params)); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ( + 0, storage->open(dir_ + "TestUniformUint8UsesQueryMetric.index", true)); + ASSERT_EQ(0, streamer->open(storage)); + + const IndexQueryMeta encoded_meta(IndexMeta::DataType::DT_INT8, + kEncodedDimension); + const auto zero_record = + EncodeUniformUint8Record(std::vector(kOriginalDimension, 0)); + const auto mid_record = + EncodeUniformUint8Record(std::vector(kOriginalDimension, 127)); + const auto zero_query = + EncodeUniformUint8Query(std::vector(kOriginalDimension, 0)); + + auto context = streamer->create_context(); + ASSERT_TRUE(context); + ASSERT_EQ(0, + streamer->add_impl(10, zero_record.data(), encoded_meta, context)); + ASSERT_EQ(0, + streamer->add_impl(20, mid_record.data(), encoded_meta, context)); + + auto brute_force_context = streamer->create_context(); + ASSERT_TRUE(brute_force_context); + brute_force_context->set_topk(2); + ASSERT_EQ(0, streamer->search_bf_impl(zero_query.data(), encoded_meta, + brute_force_context)); + ASSERT_EQ(2, brute_force_context->result().size()); + EXPECT_EQ(10, brute_force_context->result()[0].key()); + EXPECT_FLOAT_EQ(0.0f, brute_force_context->result()[0].score()); + EXPECT_EQ(20, brute_force_context->result()[1].key()); + EXPECT_FLOAT_EQ(static_cast(kOriginalDimension * 127 * 127), + brute_force_context->result()[1].score()); + + auto graph_context = streamer->create_context(); + ASSERT_TRUE(graph_context); + graph_context->set_topk(1); + ASSERT_EQ( + 0, streamer->search_impl(zero_query.data(), encoded_meta, graph_context)); + ASSERT_EQ(1, graph_context->result().size()); + EXPECT_EQ(10, graph_context->result()[0].key()); + EXPECT_FLOAT_EQ(0.0f, graph_context->result()[0].score()); + + auto primary_key_context = streamer->create_context(); + ASSERT_TRUE(primary_key_context); + primary_key_context->set_topk(2); + const std::vector> primary_keys{{10, 20}}; + ASSERT_EQ( + 0, streamer->search_bf_by_p_keys_impl(zero_query.data(), primary_keys, + encoded_meta, primary_key_context)); + ASSERT_EQ(2, primary_key_context->result().size()); + EXPECT_EQ(10, primary_key_context->result()[0].key()); + EXPECT_FLOAT_EQ(0.0f, primary_key_context->result()[0].score()); +} + // Test Vamana + INT8 quantization + rotation end-to-end TEST_F(VamanaStreamerTest, TestInt8WithRotate) { constexpr size_t kTestDim = 128; diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 26f199c49..d0570dfed 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -27,6 +27,9 @@ #include "zvec/core/framework/index_provider.h" #endif #include +#include +#include +#include "algorithm/hnsw/hnsw_params.h" #include "zvec/core/interface/index.h" #include "zvec/core/interface/index_factory.h" #include "zvec/core/interface/index_param.h" @@ -39,6 +42,25 @@ using namespace zvec::core_interface; +class ReformerInspectableHNSWIndex : public HNSWIndex { + public: + int InitForTest(const BaseIndexParam ¶m) { + return Init(param); + } + + int TransformForTest(const std::vector &query) const { + if (!reformer_) { + return zvec::core::IndexError_Uninitialized; + } + zvec::core::IndexQueryMeta input_meta( + zvec::core::IndexMeta::DataType::DT_FP32, query.size()); + zvec::core::IndexQueryMeta output_meta; + std::string output; + return reformer_->transform(query.data(), input_meta, &output, + &output_meta); + } +}; + TEST(IndexInterface, General) { constexpr uint32_t kDimension = 64; const std::string index_name{"test.index"}; @@ -190,6 +212,97 @@ TEST(IndexInterface, General) { .build()); } +TEST(IndexInterface, ReopenRestoresUniformReformer) { + constexpr size_t kDimension = 16; + struct TestCase { + QuantizerType quantizer_type; + const char *converter_name; + const char *index_name; + }; + const TestCase test_cases[] = { + {QuantizerType::kUniformUint7, "UniformUint7StreamingConverter", + "test_uniform_uint7_reopen.index"}, + {QuantizerType::kUniformUint8, "UniformUint8StreamingConverter", + "test_uniform_uint8_reopen.index"}, + }; + + for (const auto &test_case : test_cases) { + SCOPED_TRACE(test_case.converter_name); + zvec::test_util::RemoveTestFiles(test_case.index_name); + + zvec::core::IndexMeta input_meta(zvec::core::IndexMeta::DataType::DT_FP32, + kDimension); + input_meta.set_metric("SquaredEuclidean", 0, zvec::ailego::Params()); + + auto converter = + zvec::core::IndexFactory::CreateConverter(test_case.converter_name); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(input_meta, zvec::ailego::Params())); + + auto holder = std::make_shared>(kDimension); + for (uint64_t key = 0; key < 2; ++key) { + zvec::ailego::NumericalVector vector(kDimension); + for (size_t i = 0; i < kDimension; ++i) { + vector[i] = static_cast(key * kDimension + i); + } + ASSERT_TRUE(holder->emplace(key, std::move(vector))); + } + ASSERT_EQ(0, converter->train(holder)); + + zvec::ailego::Params streamer_params; + streamer_params.set(zvec::core::PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100U); + streamer_params.set(zvec::core::PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, + 16U); + streamer_params.set(zvec::core::PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, + true); + streamer_params.set(zvec::core::PARAM_HNSW_STREAMER_EF, + kDefaultHnswEfSearch); + streamer_params.set(zvec::core::PARAM_HNSW_STREAMER_USE_ID_MAP, true); + streamer_params.set(zvec::core::PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY, + false); + streamer_params.set(zvec::core::PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR, + false); + + auto streamer = zvec::core::IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(converter->meta(), streamer_params)); + + auto storage = zvec::core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(zvec::ailego::Params())); + ASSERT_EQ(0, storage->open(test_case.index_name, true)); + ASSERT_EQ(0, streamer->open(storage)); + ASSERT_EQ(0, streamer->flush(0)); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, streamer->cleanup()); + ASSERT_EQ(0, storage->close()); + + auto param = + HNSWIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithM(16) + .WithEFConstruction(100) + .WithQuantizerParam(QuantizerParam(test_case.quantizer_type)) + .Build(); + ReformerInspectableHNSWIndex index; + ASSERT_EQ(0, index.InitForTest(*param)); + + const std::vector query(kDimension, 1.0f); + ASSERT_NE(0, index.TransformForTest(query)); + ASSERT_EQ(0, index.Open(test_case.index_name, + {StorageOptions::StorageType::kMMAP, + /*create_new=*/false, /*read_only=*/true})); + EXPECT_EQ(0, index.TransformForTest(query)); + ASSERT_EQ(0, index.Close()); + + zvec::test_util::RemoveTestFiles(test_case.index_name); + } +} + TEST(IndexInterface, CopyOnWrite) { constexpr uint32_t kDimension = 64; constexpr uint32_t kNumVectors = 50; diff --git a/tests/core/interface/quantizer_param_test.cc b/tests/core/interface/quantizer_param_test.cc new file mode 100644 index 000000000..b246abbb5 --- /dev/null +++ b/tests/core/interface/quantizer_param_test.cc @@ -0,0 +1,35 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +namespace zvec::core_interface { +namespace { + +TEST(QuantizerParam, SerializesCanonicalUniformNames) { + QuantizerParam uint7(QuantizerType::kUniformUint7); + QuantizerParam uint8(QuantizerType::kUniformUint8); + EXPECT_NE(std::string::npos, uint7.SerializeToJson().find("kUniformUint7")); + EXPECT_NE(std::string::npos, uint8.SerializeToJson().find("kUniformUint8")); +} + +TEST(QuantizerParam, UsesDefaultTypeWhenJsonOmitsType) { + QuantizerParam param; + ASSERT_TRUE(param.DeserializeFromJson("{}")); + EXPECT_EQ(QuantizerType::kNone, param.type); +} + +} // namespace +} // namespace zvec::core_interface diff --git a/tests/core/metric/uniform_uint7_metric_test.cc b/tests/core/metric/uniform_uint7_metric_test.cc new file mode 100644 index 000000000..7dbd30761 --- /dev/null +++ b/tests/core/metric/uniform_uint7_metric_test.cc @@ -0,0 +1,76 @@ +// 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 + +namespace zvec::core { +namespace { + +TEST(UniformUint7Metric, UsesCanonicalParamsAndComputesDistance) { + auto metric = IndexFactory::CreateMetric("UniformUint7"); + ASSERT_TRUE(metric); + + ailego::Params params; + params.set("proxima.uniform_uint7.metric.origin_metric_name", + std::string("SquaredEuclidean")); + constexpr size_t kDimension = 4; + ASSERT_EQ(0, metric->init(IndexMeta(IndexMeta::DataType::DT_INT8, kDimension), + params)); + + std::string metric_name; + EXPECT_TRUE(metric->params().get( + "proxima.uniform_uint7.metric.origin_metric_name", &metric_name)); + EXPECT_EQ("SquaredEuclidean", metric_name); + + const int8_t lhs[kDimension] = {0, 1, 2, 127}; + const int8_t rhs[kDimension] = {0, 2, 4, 120}; + float distance = -1.0f; + metric->distance()(lhs, rhs, kDimension, &distance); + EXPECT_FLOAT_EQ(54.0f, distance); + + const IndexMeta matching_meta(IndexMeta::DataType::DT_INT8, kDimension); + const IndexMeta wrong_dimension_meta(IndexMeta::DataType::DT_INT8, + kDimension + 1); + EXPECT_TRUE(metric->is_matched(matching_meta)); + EXPECT_FALSE(metric->is_matched(wrong_dimension_meta)); + EXPECT_TRUE(metric->is_matched( + matching_meta, IndexQueryMeta(IndexMeta::DataType::DT_INT8, kDimension))); + EXPECT_FALSE(metric->is_matched( + matching_meta, + IndexQueryMeta(IndexMeta::DataType::DT_INT8, kDimension + 1))); + EXPECT_FALSE(metric->is_matched( + wrong_dimension_meta, + IndexQueryMeta(IndexMeta::DataType::DT_INT8, kDimension + 1))); +} + +TEST(UniformUint7Metric, RejectsInvalidDimension) { + ailego::Params params; + params.set("proxima.uniform_uint7.metric.origin_metric_name", + std::string("SquaredEuclidean")); + for (const uint32_t dimension : + {0U, static_cast(MAX_DIMENSION + 1)}) { + auto metric = IndexFactory::CreateMetric("UniformUint7"); + ASSERT_TRUE(metric); + EXPECT_EQ(IndexError_InvalidArgument, + metric->init(IndexMeta(IndexMeta::DataType::DT_INT8, dimension), + params)); + } +} + +} // namespace +} // namespace zvec::core diff --git a/tests/core/metric/uniform_uint8_metric_test.cc b/tests/core/metric/uniform_uint8_metric_test.cc new file mode 100644 index 000000000..f21b7536a --- /dev/null +++ b/tests/core/metric/uniform_uint8_metric_test.cc @@ -0,0 +1,410 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "metric/metric_params.h" + +namespace zvec::core { +namespace { + +constexpr size_t kTailBytes = sizeof(uint32_t); + +std::vector EncodeRecord(const std::vector &codes) { + std::vector encoded(codes.size() + kTailBytes, 0); + int64_t sum_squared = 0; + for (size_t i = 0; i < codes.size(); ++i) { + encoded[i] = static_cast(static_cast(codes[i]) - 128); + sum_squared += static_cast(codes[i]) * codes[i]; + } + const uint32_t tail = static_cast(sum_squared); + std::memcpy(encoded.data() + codes.size(), &tail, sizeof(tail)); + return encoded; +} + +std::vector EncodeQuery(const std::vector &codes) { + std::vector encoded(codes.size() + kTailBytes, 0); + int64_t sum_squared = 0; + auto *bytes = reinterpret_cast(encoded.data()); + for (size_t i = 0; i < codes.size(); ++i) { + bytes[i] = codes[i]; + sum_squared += static_cast(codes[i]) * codes[i]; + } + const uint32_t tail = static_cast(sum_squared); + std::memcpy(encoded.data() + codes.size(), &tail, sizeof(tail)); + return encoded; +} + +uint32_t ReadTail(const std::vector &encoded, + size_t original_dimension) { + uint32_t tail = 0; + std::memcpy(&tail, encoded.data() + original_dimension, sizeof(tail)); + return tail; +} + +int32_t ReadQueryCorrection(const std::vector &encoded, + size_t original_dimension) { + int32_t correction = 0; + std::memcpy(&correction, encoded.data() + original_dimension, + sizeof(correction)); + return correction; +} + +int64_t SquaredL2(const std::vector &lhs, + const std::vector &rhs) { + int64_t result = 0; + for (size_t i = 0; i < lhs.size(); ++i) { + const int difference = static_cast(lhs[i]) - rhs[i]; + result += static_cast(difference) * difference; + } + return result; +} + +std::vector PrepareQuery(const IndexMetric::Pointer &query_metric, + const std::vector &query) { + auto prepared = query; + const auto preprocess = query_metric->get_query_preprocess_func(); + EXPECT_TRUE(preprocess); + if (preprocess) { + preprocess(prepared.data(), prepared.size()); + } + return prepared; +} + +IndexMetric::Pointer CreateMetric(size_t original_dimension) { + auto metric = IndexFactory::CreateMetric("UniformUint8"); + if (!metric) { + return nullptr; + } + IndexMeta meta(IndexMeta::DataType::DT_INT8, original_dimension + kTailBytes); + ailego::Params params; + params.set(UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME, + std::string("SquaredEuclidean")); + return metric->init(meta, params) == 0 ? metric : nullptr; +} + +TEST(UniformUint8Metric, UsesExactBuildAndQueryDistance) { + const std::vector first_codes{0, 1, 127, 128, 254, 255}; + const std::vector second_codes{255, 128, 127, 1, 0, 254}; + const std::vector query_codes{17, 255, 3, 200, 128, 0}; + const auto first = EncodeRecord(first_codes); + const auto second = EncodeRecord(second_codes); + const auto query = EncodeQuery(query_codes); + + auto metric = CreateMetric(first_codes.size()); + ASSERT_TRUE(metric); + ASSERT_TRUE(metric->distance()); + ASSERT_TRUE(metric->batch_distance()); + ASSERT_TRUE(metric->get_query_preprocess_func()); + auto query_metric = metric->query_metric(); + ASSERT_TRUE(query_metric); + ASSERT_TRUE(query_metric->distance()); + + float distance = 0.0f; + metric->distance()(first.data(), second.data(), first.size(), &distance); + EXPECT_FLOAT_EQ(static_cast(SquaredL2(first_codes, second_codes)), + distance); + + auto prepared_stored_query = second; + metric->get_query_preprocess_func()(prepared_stored_query.data(), + prepared_stored_query.size()); + const void *stored_vectors[] = {first.data()}; + metric->batch_distance()(stored_vectors, prepared_stored_query.data(), 1, + first.size(), &distance); + EXPECT_FLOAT_EQ(static_cast(SquaredL2(first_codes, second_codes)), + distance); + + query_metric->distance()(first.data(), query.data(), first.size(), &distance); + EXPECT_FLOAT_EQ(static_cast(SquaredL2(first_codes, query_codes)), + distance); + + const auto prepared_query = PrepareQuery(query_metric, query); + int64_t expected_correction = 0; + for (uint8_t code : query_codes) { + expected_correction += + static_cast(code) * code - 256 * static_cast(code); + } + EXPECT_EQ(expected_correction, + ReadQueryCorrection(prepared_query, query_codes.size())); + query_metric->distance()(first.data(), prepared_query.data(), first.size(), + &distance); + EXPECT_FLOAT_EQ(static_cast(SquaredL2(first_codes, query_codes)), + distance); +} + +TEST(UniformUint8Metric, QueryBatchMatchesScalarAcrossKernelBoundaries) { + constexpr size_t kVectorCount = 7; + for (const size_t dimension : + {1UL, 15UL, 16UL, 31UL, 63UL, 64UL, 65UL, 127UL, 128UL, 129UL, 1024UL}) { + std::vector query_codes(dimension); + std::vector> vector_codes( + kVectorCount, std::vector(dimension)); + for (size_t d = 0; d < dimension; ++d) { + query_codes[d] = static_cast((d * 131 + 17) & 0xff); + for (size_t i = 0; i < kVectorCount; ++i) { + vector_codes[i][d] = + static_cast((d * (29 + i * 18) + i * 53) & 0xff); + } + } + + const auto query = EncodeQuery(query_codes); + std::vector> vectors; + std::vector vector_pointers; + vectors.reserve(kVectorCount); + vector_pointers.reserve(kVectorCount); + for (const auto &codes : vector_codes) { + vectors.push_back(EncodeRecord(codes)); + vector_pointers.push_back(vectors.back().data()); + } + + auto metric = CreateMetric(dimension); + ASSERT_TRUE(metric); + auto query_metric = metric->query_metric(); + ASSERT_TRUE(query_metric); + ASSERT_TRUE(query_metric->batch_distance()); + // Streamer contexts retain the build metric while swapping in the query + // distance functions, so preprocessing must also be available there. + const auto prepared_query = PrepareQuery(metric, query); + + std::vector distances(kVectorCount); + query_metric->batch_distance()(vector_pointers.data(), + prepared_query.data(), kVectorCount, + dimension + kTailBytes, distances.data()); + for (size_t i = 0; i < kVectorCount; ++i) { + EXPECT_FLOAT_EQ( + static_cast(SquaredL2(vector_codes[i], query_codes)), + distances[i]) + << "dimension=" << dimension << ", vector=" << i; + } + } +} + +TEST(UniformUint8Metric, QueryPreprocessIsIdempotent) { + constexpr size_t kDimension = MAX_DIMENSION; + const std::vector query_codes(kDimension, uint8_t{255}); + auto query = EncodeQuery(query_codes); + + auto metric = CreateMetric(kDimension); + ASSERT_TRUE(metric); + const auto preprocess = metric->get_query_preprocess_func(); + ASSERT_TRUE(preprocess); + + preprocess(query.data(), query.size()); + const auto once_preprocessed = query; + EXPECT_EQ(-static_cast(kDimension) * 255, + ReadQueryCorrection(query, kDimension)); + + preprocess(query.data(), query.size()); + EXPECT_EQ(once_preprocessed, query); +} + +TEST(UniformUint8Metric, TurboBatchCallWritesEveryDistanceWhenAvailable) { + constexpr size_t kVectorCount = 3; + const std::vector query_codes{17, 255, 3, 200, 128, 0}; + const std::vector> record_codes{ + {0, 1, 127, 128, 254, 255}, + {255, 128, 127, 1, 0, 254}, + {33, 66, 99, 132, 165, 198}, + }; + const auto query = EncodeQuery(query_codes); + + std::vector> records; + std::vector record_pointers; + records.reserve(kVectorCount); + record_pointers.reserve(kVectorCount); + for (const auto &codes : record_codes) { + records.push_back(EncodeRecord(codes)); + record_pointers.push_back(records.back().data()); + } + + std::vector distances(kVectorCount, + std::numeric_limits::quiet_NaN()); + auto batch_distance = zvec::turbo::get_batch_distance_func( + zvec::turbo::MetricType::kSquaredEuclidean, zvec::turbo::DataType::kInt8, + zvec::turbo::QuantizeType::kUniformUint8); + auto preprocess = zvec::turbo::get_query_preprocess_func( + zvec::turbo::MetricType::kSquaredEuclidean, zvec::turbo::DataType::kInt8, + zvec::turbo::QuantizeType::kUniformUint8); + if (!batch_distance || !preprocess) { + GTEST_SKIP() << "AVX512-VNNI is not available on this CPU"; + } + auto prepared_query = query; + preprocess(prepared_query.data(), prepared_query.size()); + batch_distance(record_pointers.data(), prepared_query.data(), kVectorCount, + query_codes.size() + kTailBytes, distances.data()); + + for (size_t i = 0; i < kVectorCount; ++i) { + EXPECT_FLOAT_EQ(static_cast(SquaredL2(record_codes[i], query_codes)), + distances[i]); + } +} + +TEST(UniformUint8Metric, + ExactQueryDistanceSupportsUint32RangeAtMaximumDimension) { + constexpr size_t kDimension = MAX_DIMENSION; + constexpr size_t kVectorCount = 4; + + auto metric = CreateMetric(kDimension); + ASSERT_TRUE(metric); + auto query_metric = metric->query_metric(); + ASSERT_TRUE(query_metric); + ASSERT_TRUE(query_metric->distance()); + ASSERT_TRUE(query_metric->batch_distance()); + + const auto verify = [&](uint8_t record_code, uint8_t query_code) { + const std::vector record_codes(kDimension, record_code); + const std::vector query_codes(kDimension, query_code); + const auto record = EncodeRecord(record_codes); + const auto query = EncodeQuery(query_codes); + const int64_t expected = SquaredL2(record_codes, query_codes); + EXPECT_EQ(int64_t{kDimension} * 255 * 255, expected); + ASSERT_GT(expected, (std::numeric_limits::max)()); + ASSERT_LE(static_cast(expected), + (std::numeric_limits::max)()); + + float scalar_distance = 0.0f; + query_metric->distance()(record.data(), query.data(), + kDimension + kTailBytes, &scalar_distance); + EXPECT_FLOAT_EQ(static_cast(expected), scalar_distance); + + const auto prepared_query = PrepareQuery(query_metric, query); + const void *records[kVectorCount] = {record.data(), record.data(), + record.data(), record.data()}; + float batch_distances[kVectorCount] = {}; + query_metric->batch_distance()(records, prepared_query.data(), kVectorCount, + kDimension + kTailBytes, batch_distances); + for (float distance : batch_distances) { + EXPECT_FLOAT_EQ(static_cast(expected), distance); + } + }; + + verify(/*record_code=*/255, /*query_code=*/0); + verify(/*record_code=*/0, /*query_code=*/255); + + const auto full_range_record = + EncodeRecord(std::vector(kDimension, uint8_t{255})); + EXPECT_EQ(uint64_t{kDimension} * 255 * 255, + ReadTail(full_range_record, kDimension)); +} + +TEST(UniformUint8Metric, ExactQueryDistanceFallsBackAboveTurboDimensionLimit) { + constexpr size_t kDimension = MAX_DIMENSION + 4096; + constexpr size_t kVectorCount = 4; + + auto metric = CreateMetric(kDimension); + ASSERT_TRUE(metric); + auto query_metric = metric->query_metric(); + ASSERT_TRUE(query_metric); + ASSERT_TRUE(query_metric->distance()); + ASSERT_TRUE(query_metric->batch_distance()); + + const auto verify = [&](uint8_t record_code, uint8_t query_code) { + const std::vector record_codes(kDimension, record_code); + const std::vector query_codes(kDimension, query_code); + const auto record = EncodeRecord(record_codes); + const auto query = EncodeQuery(query_codes); + const int64_t expected = SquaredL2(record_codes, query_codes); + ASSERT_GT(expected, (std::numeric_limits::max)()); + + float scalar_distance = 0.0f; + query_metric->distance()(record.data(), query.data(), + kDimension + kTailBytes, &scalar_distance); + EXPECT_FLOAT_EQ(static_cast(expected), scalar_distance); + + const auto prepared_query = PrepareQuery(query_metric, query); + const void *records[kVectorCount] = {record.data(), record.data(), + record.data(), record.data()}; + float distances[kVectorCount] = {}; + query_metric->batch_distance()(records, prepared_query.data(), kVectorCount, + kDimension + kTailBytes, distances); + for (float distance : distances) { + EXPECT_FLOAT_EQ(static_cast(expected), distance); + } + }; + + verify(/*record_code=*/0, /*query_code=*/255); + verify(/*record_code=*/255, /*query_code=*/0); +} + +TEST(UniformUint8Metric, RejectsEncodedDimensionWithoutVectorData) { + for (const uint32_t encoded_dimension : + {0U, static_cast(kTailBytes)}) { + auto metric = IndexFactory::CreateMetric("UniformUint8"); + ASSERT_TRUE(metric); + IndexMeta meta(IndexMeta::DataType::DT_INT8, encoded_dimension); + ailego::Params params; + params.set(UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME, + std::string("SquaredEuclidean")); + EXPECT_NE(0, metric->init(meta, params)) + << "encoded_dimension=" << encoded_dimension; + } +} + +TEST(UniformUint8Metric, BuildDistanceIsExactForFullRangeAndLargeDimensions) { + std::mt19937 generator(20260716); + std::uniform_int_distribution byte_distribution(0, 255); + + for (const size_t dimension : {1UL, 31UL, 32UL, 33UL, 127UL, 128UL, 129UL, + 1024UL, 65536UL, 1048577UL}) { + std::vector lhs_codes(dimension); + std::vector rhs_codes(dimension); + for (size_t i = 0; i < dimension; ++i) { + lhs_codes[i] = static_cast(byte_distribution(generator)); + rhs_codes[i] = static_cast(byte_distribution(generator)); + } + auto lhs = EncodeRecord(lhs_codes); + auto rhs = EncodeRecord(rhs_codes); + + auto metric = CreateMetric(dimension); + ASSERT_TRUE(metric); + float distance = 0.0f; + metric->distance()(lhs.data(), rhs.data(), dimension + kTailBytes, + &distance); + EXPECT_FLOAT_EQ(static_cast(SquaredL2(lhs_codes, rhs_codes)), + distance) + << "dimension=" << dimension; + + std::fill(lhs_codes.begin(), lhs_codes.end(), uint8_t{0}); + std::fill(rhs_codes.begin(), rhs_codes.end(), uint8_t{255}); + lhs = EncodeRecord(lhs_codes); + rhs = EncodeRecord(rhs_codes); + metric->distance()(lhs.data(), rhs.data(), dimension + kTailBytes, + &distance); + EXPECT_FLOAT_EQ(static_cast(SquaredL2(lhs_codes, rhs_codes)), + distance) + << "extreme dimension=" << dimension; + } +} + +TEST(UniformUint8Metric, RejectsUnsupportedOriginMetric) { + auto metric = IndexFactory::CreateMetric("UniformUint8"); + ASSERT_TRUE(metric); + IndexMeta meta(IndexMeta::DataType::DT_INT8, 16 + kTailBytes); + ailego::Params params; + params.set(UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME, + std::string("InnerProduct")); + EXPECT_NE(0, metric->init(meta, params)); +} + +} // namespace +} // namespace zvec::core diff --git a/tests/core/quantizer/uniform_int8_reformer_test.cc b/tests/core/quantizer/uniform_uint7_reformer_test.cc similarity index 65% rename from tests/core/quantizer/uniform_int8_reformer_test.cc rename to tests/core/quantizer/uniform_uint7_reformer_test.cc index ae47a88f0..0db8f4ac3 100644 --- a/tests/core/quantizer/uniform_int8_reformer_test.cc +++ b/tests/core/quantizer/uniform_uint7_reformer_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,13 +24,24 @@ #include #include "zvec/core/framework/index_factory.h" #include "zvec/core/framework/index_holder.h" +#include "zvec/core/interface/index_param.h" using namespace zvec::core; +namespace { + +float ReadFloat(const std::string &data, size_t index) { + float value = 0.0f; + std::memcpy(&value, data.data() + index * sizeof(value), sizeof(value)); + return value; +} + +} // namespace + // --------------------------------------------------------------------------- -// UniformInt8 Converter + Reformer: General (MultiPassHolder, uniform dist) +// UniformUint7 Converter + Reformer: General (MultiPassHolder, uniform dist) // --------------------------------------------------------------------------- -TEST(UniformInt8Reformer, General) { +TEST(UniformUint7Reformer, General) { std::mt19937 gen(42); std::uniform_real_distribution dist(0.0f, 1.0f); @@ -39,7 +52,7 @@ TEST(UniformInt8Reformer, General) { meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); auto converter = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter); ASSERT_EQ(0u, converter->init(meta, zvec::ailego::Params())); @@ -82,7 +95,7 @@ TEST(UniformInt8Reformer, General) { } // Create reformer from converter's trained params - auto reformer = IndexFactory::CreateReformer("UniformInt8StreamingReformer"); + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); ASSERT_TRUE(reformer); ASSERT_EQ(0u, reformer->init(converter->meta().reformer_params())); @@ -139,10 +152,59 @@ TEST(UniformInt8Reformer, General) { } } +TEST(UniformUint7Reformer, UsesNearestEvenRoundingAcrossSimdBoundary) { + constexpr size_t kDimension = 20; + const float input[kDimension] = { + -1.5f, -0.5f, 0.5f, 1.5f, 2.5f, 3.5f, 4.5f, 5.5f, 6.5f, 7.5f, + 8.5f, 9.5f, 10.5f, 11.5f, 126.5f, 127.5f, 0.5f, 1.5f, 2.5f, 3.5f, + }; + const int8_t expected[kDimension] = { + 0, 0, 0, 2, 2, 4, 4, 6, 6, 8, 8, 10, 10, 12, 126, 127, 0, 2, 2, 4, + }; + + zvec::ailego::Params params; + params.set("uniform_uint7.reformer.scale", 1.0f); + params.set("uniform_uint7.reformer.bias", 0.0f); + + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(params)); + + std::string output; + IndexQueryMeta output_meta; + ASSERT_EQ(0, + reformer->transform( + input, IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension), + &output, &output_meta)); + ASSERT_EQ(kDimension, output.size()); + EXPECT_EQ(0, std::memcmp(expected, output.data(), kDimension)); + + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + auto converter = + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0, converter->init(meta, params)); + + auto holder = + std::make_shared>( + kDimension); + zvec::ailego::NumericalVector vector(kDimension); + std::copy(input, input + kDimension, vector.data()); + holder->emplace(1, vector); + ASSERT_EQ(0, converter->transform(holder)); + + auto result = converter->result(); + ASSERT_TRUE(result); + auto iterator = result->create_iterator(); + ASSERT_TRUE(iterator); + ASSERT_TRUE(iterator->is_valid()); + EXPECT_EQ(0, std::memcmp(expected, iterator->data(), kDimension)); +} + // --------------------------------------------------------------------------- // OnePassHolder: verify converter works with single-pass holders // --------------------------------------------------------------------------- -TEST(UniformInt8Reformer, OnePassHolder) { +TEST(UniformUint7Reformer, OnePassHolder) { std::mt19937 gen(123); std::normal_distribution dist(5.0f, 2.0f); @@ -153,7 +215,7 @@ TEST(UniformInt8Reformer, OnePassHolder) { meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); auto converter = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter); ASSERT_EQ(0u, converter->init(meta, zvec::ailego::Params())); @@ -180,7 +242,7 @@ TEST(UniformInt8Reformer, OnePassHolder) { EXPECT_EQ(IndexMeta::DataType::DT_INT8, holder2->data_type()); EXPECT_EQ(DIMENSION, holder2->dimension()); - auto reformer = IndexFactory::CreateReformer("UniformInt8StreamingReformer"); + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); ASSERT_TRUE(reformer); ASSERT_EQ(0u, reformer->init(converter->meta().reformer_params())); @@ -204,9 +266,9 @@ TEST(UniformInt8Reformer, OnePassHolder) { } // --------------------------------------------------------------------------- -// TrainedParams: verify scale/bias are persisted correctly after train +// TrainedParams: verify scale/bias are published in metadata after train // --------------------------------------------------------------------------- -TEST(UniformInt8Reformer, TrainedParams) { +TEST(UniformUint7Reformer, TrainedParams) { std::mt19937 gen(99); std::uniform_real_distribution dist(-3.0f, 7.0f); @@ -217,7 +279,7 @@ TEST(UniformInt8Reformer, TrainedParams) { meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); auto converter = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter); ASSERT_EQ(0u, converter->init(meta, zvec::ailego::Params())); @@ -238,8 +300,8 @@ TEST(UniformInt8Reformer, TrainedParams) { // Verify reformer params contain scale and bias auto reformer_params = converter->meta().reformer_params(); float scale = 0.0f, bias = 0.0f; - EXPECT_TRUE(reformer_params.get("uniform_int8.reformer.scale", &scale)); - EXPECT_TRUE(reformer_params.get("uniform_int8.reformer.bias", &bias)); + EXPECT_TRUE(reformer_params.get("uniform_uint7.reformer.scale", &scale)); + EXPECT_TRUE(reformer_params.get("uniform_uint7.reformer.bias", &bias)); EXPECT_GT(scale, 0.0f); EXPECT_TRUE(std::isfinite(scale)); EXPECT_TRUE(std::isfinite(bias)); @@ -247,20 +309,20 @@ TEST(UniformInt8Reformer, TrainedParams) { // Verify converter params also contain scale/bias (for persistence) auto conv_params = converter->meta().converter_params(); float conv_scale = 0.0f, conv_bias = 0.0f; - EXPECT_TRUE(conv_params.get("uniform_int8.reformer.scale", &conv_scale)); - EXPECT_TRUE(conv_params.get("uniform_int8.reformer.bias", &conv_bias)); + EXPECT_TRUE(conv_params.get("uniform_uint7.reformer.scale", &conv_scale)); + EXPECT_TRUE(conv_params.get("uniform_uint7.reformer.bias", &conv_bias)); EXPECT_FLOAT_EQ(scale, conv_scale); EXPECT_FLOAT_EQ(bias, conv_bias); // Verify meta reflects the correct reformer and metric - EXPECT_EQ("UniformInt8StreamingReformer", converter->meta().reformer_name()); - EXPECT_EQ("UniformInt8", converter->meta().metric_name()); + EXPECT_EQ("UniformUint7StreamingReformer", converter->meta().reformer_name()); + EXPECT_EQ("UniformUint7", converter->meta().metric_name()); } // --------------------------------------------------------------------------- // Revert: verify int8 → float dequantization round-trip quality // --------------------------------------------------------------------------- -TEST(UniformInt8Reformer, Revert) { +TEST(UniformUint7Reformer, Revert) { std::mt19937 gen(77); std::uniform_real_distribution dist(0.0f, 10.0f); @@ -271,7 +333,7 @@ TEST(UniformInt8Reformer, Revert) { meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); auto converter = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter); ASSERT_EQ(0u, converter->init(meta, zvec::ailego::Params())); @@ -288,7 +350,7 @@ TEST(UniformInt8Reformer, Revert) { ASSERT_EQ(0u, IndexConverter::TrainAndTransform(converter, holder)); - auto reformer = IndexFactory::CreateReformer("UniformInt8StreamingReformer"); + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); ASSERT_TRUE(reformer); ASSERT_EQ(0u, reformer->init(converter->meta().reformer_params())); @@ -307,17 +369,15 @@ TEST(UniformInt8Reformer, Revert) { ASSERT_EQ(0, reformer->revert(quantized_buf.data(), qmeta, &reverted_buf)); - const float *reverted = - reinterpret_cast(reverted_buf.data()); - // Quantization error should be bounded by step_size / 2 // step_size ≈ range / 127 float range = 10.0f; // approximate float max_error = range / 127.0f; for (size_t d = 0; d < DIMENSION; ++d) { - EXPECT_NEAR(original[d], reverted[d], max_error * 1.5f) + const float reverted = ReadFloat(reverted_buf, d); + EXPECT_NEAR(original[d], reverted, max_error * 1.5f) << "dim=" << d << " original=" << original[d] - << " reverted=" << reverted[d]; + << " reverted=" << reverted; } } } @@ -325,7 +385,7 @@ TEST(UniformInt8Reformer, Revert) { // --------------------------------------------------------------------------- // Normalize: verify score rescaling from int8 L2 to float L2 // --------------------------------------------------------------------------- -TEST(UniformInt8Reformer, Normalize) { +TEST(UniformUint7Reformer, Normalize) { const size_t COUNT = 1000; const size_t DIMENSION = 32; @@ -336,7 +396,7 @@ TEST(UniformInt8Reformer, Normalize) { meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); auto converter = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter); ASSERT_EQ(0u, converter->init(meta, zvec::ailego::Params())); @@ -355,9 +415,9 @@ TEST(UniformInt8Reformer, Normalize) { auto reformer_params = converter->meta().reformer_params(); float scale = 0.0f; - ASSERT_TRUE(reformer_params.get("uniform_int8.reformer.scale", &scale)); + ASSERT_TRUE(reformer_params.get("uniform_uint7.reformer.scale", &scale)); - auto reformer = IndexFactory::CreateReformer("UniformInt8StreamingReformer"); + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); ASSERT_TRUE(reformer); ASSERT_EQ(0u, reformer->init(reformer_params)); @@ -379,10 +439,10 @@ TEST(UniformInt8Reformer, Normalize) { } // --------------------------------------------------------------------------- -// InitConverterWithTrainedParams: simulate the search-only path where -// scale/bias come from persisted converter params (no re-train needed) +// InitConverterWithTrainedParams: verify a converter can reuse scale/bias from +// previously trained converter params without retraining // --------------------------------------------------------------------------- -TEST(UniformInt8Reformer, InitConverterWithTrainedParams) { +TEST(UniformUint7Reformer, InitConverterWithTrainedParams) { std::mt19937 gen(42); std::uniform_real_distribution dist(0.0f, 1.0f); @@ -394,7 +454,7 @@ TEST(UniformInt8Reformer, InitConverterWithTrainedParams) { // First pass: train to get params auto converter = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter); ASSERT_EQ(0u, converter->init(meta, zvec::ailego::Params())); @@ -415,7 +475,7 @@ TEST(UniformInt8Reformer, InitConverterWithTrainedParams) { // Second pass: create a new converter with trained params (skip train) auto converter2 = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter2); ASSERT_EQ(0, converter2->init(meta, converter_params)); ASSERT_EQ(0, converter2->transform(holder)); @@ -430,8 +490,8 @@ TEST(UniformInt8Reformer, InitConverterWithTrainedParams) { EXPECT_EQ(IndexMeta::DataType::DT_INT8, holder2->data_type()); EXPECT_EQ(DIMENSION, holder2->dimension()); - // Verify reformer with persisted params produces same results - auto reformer = IndexFactory::CreateReformer("UniformInt8StreamingReformer"); + // Verify a reformer initialized with the trained params produces same results + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); ASSERT_TRUE(reformer); ASSERT_EQ(0u, reformer->init(reformer_params)); @@ -467,7 +527,7 @@ TEST(UniformInt8Reformer, InitConverterWithTrainedParams) { // LosslessIntegerFastPath: when all training values are integers within // [0, 127], scale should be 1.0 for exact mapping // --------------------------------------------------------------------------- -TEST(UniformInt8Reformer, LosslessIntegerFastPath) { +TEST(UniformUint7Reformer, LosslessIntegerFastPath) { const size_t COUNT = 100; const size_t DIMENSION = 8; @@ -475,7 +535,7 @@ TEST(UniformInt8Reformer, LosslessIntegerFastPath) { meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); auto converter = - IndexFactory::CreateConverter("UniformInt8StreamingConverter"); + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); ASSERT_TRUE(converter); ASSERT_EQ(0u, converter->init(meta, zvec::ailego::Params())); @@ -499,11 +559,11 @@ TEST(UniformInt8Reformer, LosslessIntegerFastPath) { // scale should be 1.0 for lossless integer path auto reformer_params = converter->meta().reformer_params(); float scale = 0.0f; - ASSERT_TRUE(reformer_params.get("uniform_int8.reformer.scale", &scale)); + ASSERT_TRUE(reformer_params.get("uniform_uint7.reformer.scale", &scale)); EXPECT_FLOAT_EQ(1.0f, scale); // Verify exact round-trip for integer values - auto reformer = IndexFactory::CreateReformer("UniformInt8StreamingReformer"); + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); ASSERT_TRUE(reformer); ASSERT_EQ(0u, reformer->init(reformer_params)); @@ -530,10 +590,112 @@ TEST(UniformInt8Reformer, LosslessIntegerFastPath) { // Revert should give exact values back ASSERT_EQ(0, reformer->revert(quantized_buf.data(), qmeta, &reverted_buf)); - const float *reverted = - reinterpret_cast(reverted_buf.data()); for (size_t d = 0; d < DIMENSION; ++d) { - EXPECT_FLOAT_EQ(original[d], reverted[d]) << "dim=" << d; + EXPECT_FLOAT_EQ(original[d], ReadFloat(reverted_buf, d)) << "dim=" << d; } } } + +TEST(UniformUint7Converter, RejectsMismatchedTrainingHolder) { + constexpr size_t kDimension = 8; + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + auto converter = + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0, converter->init(meta, zvec::ailego::Params())); + + auto wrong_type = + std::make_shared>( + kDimension); + EXPECT_EQ(IndexError_Mismatch, converter->train(wrong_type)); + + auto wrong_dimension = + std::make_shared>( + kDimension + 1); + EXPECT_EQ(IndexError_Mismatch, converter->train(wrong_dimension)); +} + +TEST(UniformUint7Converter, RejectsInvalidDimension) { + for (const uint32_t dimension : {0U, uint32_t{MAX_DIMENSION + 1}}) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, dimension); + auto converter = + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); + ASSERT_TRUE(converter); + EXPECT_EQ(IndexError_InvalidArgument, + converter->init(meta, zvec::ailego::Params())) + << "dimension=" << dimension; + } +} + +TEST(UniformUint7Reformer, RejectsInvalidTransformAndRevertArguments) { + constexpr size_t kDimension = 4; + zvec::ailego::Params params; + params.set("uniform_uint7.reformer.scale", 1.0f); + params.set("uniform_uint7.reformer.bias", 0.0f); + + auto reformer = IndexFactory::CreateReformer("UniformUint7StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(params)); + + const std::vector input(kDimension, 1.0f); + const IndexQueryMeta input_meta(IndexMeta::DataType::DT_FP32, kDimension); + std::string output; + IndexQueryMeta output_meta; + + EXPECT_EQ(IndexError_InvalidArgument, + reformer->transform(nullptr, input_meta, &output, &output_meta)); + EXPECT_EQ( + IndexError_InvalidArgument, + reformer->transform(input.data(), input_meta, nullptr, &output_meta)); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->transform(input.data(), input_meta, &output, nullptr)); + EXPECT_EQ( + IndexError_InvalidArgument, + reformer->transform(input.data(), input_meta, 0, &output, &output_meta)); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->transform(input.data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, 0), + &output, &output_meta)); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->transform( + input.data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, MAX_DIMENSION + 1), + &output, &output_meta)); + + const std::vector encoded(kDimension, 0); + const IndexQueryMeta encoded_meta(IndexMeta::DataType::DT_INT8, kDimension); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->revert(nullptr, encoded_meta, &output)); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->revert(encoded.data(), encoded_meta, nullptr)); + EXPECT_EQ( + IndexError_InvalidArgument, + reformer->revert(encoded.data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension), + &output)); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->revert(encoded.data(), + IndexQueryMeta(IndexMeta::DataType::DT_INT8, 0), + &output)); +} + +TEST(UniformUint7Converter, RejectsNonFiniteTrainingRange) { + constexpr size_t kDimension = 1; + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + auto converter = + IndexFactory::CreateConverter("UniformUint7StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0, converter->init(meta, zvec::ailego::Params())); + + auto holder = + std::make_shared>( + kDimension); + zvec::ailego::NumericalVector low(kDimension); + low[0] = std::numeric_limits::lowest(); + holder->emplace(1, low); + zvec::ailego::NumericalVector high(kDimension); + high[0] = std::numeric_limits::max(); + holder->emplace(2, high); + + EXPECT_EQ(IndexError_InvalidArgument, converter->train(holder)); +} diff --git a/tests/core/quantizer/uniform_uint8_reformer_test.cc b/tests/core/quantizer/uniform_uint8_reformer_test.cc new file mode 100644 index 000000000..da0a62d3d --- /dev/null +++ b/tests/core/quantizer/uniform_uint8_reformer_test.cc @@ -0,0 +1,342 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec::core { +namespace { + +constexpr size_t kDimension = 6; +constexpr size_t kEncodedDimension = kDimension + sizeof(uint32_t); + +uint32_t SumSquared(const std::vector &vector) { + int64_t result = 0; + for (const float value : vector) { + const int code = static_cast(value); + result += code * code; + } + return static_cast(result); +} + +uint32_t ReadTail(const void *data, size_t dimension = kDimension) { + uint32_t tail = 0; + std::memcpy(&tail, static_cast(data) + dimension, + sizeof(tail)); + return tail; +} + +float ReadFloat(const std::string &data, size_t index) { + float value = 0.0f; + std::memcpy(&value, data.data() + index * sizeof(value), sizeof(value)); + return value; +} + +IndexHolder::Pointer MakeHolder( + const std::vector> &vectors) { + auto holder = + std::make_shared>( + kDimension); + for (size_t i = 0; i < vectors.size(); ++i) { + ailego::NumericalVector vector(kDimension); + for (size_t j = 0; j < kDimension; ++j) { + vector[j] = vectors[i][j]; + } + holder->emplace(i + 1, vector); + } + return holder; +} + +TEST(UniformUint8Reformer, EncodesFullRangeAndStoresTrainingParamsInMeta) { + const std::vector> vectors{ + {0.0f, 1.0f, 127.0f, 128.0f, 254.0f, 255.0f}, + {255.0f, 200.0f, 128.0f, 127.0f, 1.0f, 0.0f}, + }; + auto holder = MakeHolder(vectors); + + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto converter = + IndexFactory::CreateConverter("UniformUint8StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0, converter->init(meta, ailego::Params())); + ASSERT_EQ(0, IndexConverter::TrainAndTransform(converter, holder)); + + EXPECT_EQ("UniformUint8", converter->meta().metric_name()); + EXPECT_EQ("UniformUint8StreamingReformer", converter->meta().reformer_name()); + EXPECT_EQ(IndexMeta::DataType::DT_INT8, converter->meta().data_type()); + EXPECT_EQ(kEncodedDimension, converter->meta().dimension()); + float scale = 0.0f; + float bias = 0.0f; + EXPECT_TRUE(converter->meta().converter_params().get( + "uniform_uint8.reformer.scale", &scale)); + EXPECT_TRUE(converter->meta().converter_params().get( + "uniform_uint8.reformer.bias", &bias)); + EXPECT_FLOAT_EQ(1.0f, scale); + EXPECT_FLOAT_EQ(0.0f, bias); + + auto converted_holder = converter->result(); + ASSERT_TRUE(converted_holder); + ASSERT_EQ(kEncodedDimension, converted_holder->dimension()); + ASSERT_EQ(kEncodedDimension, converted_holder->element_size()); + auto iterator = converted_holder->create_iterator(); + ASSERT_TRUE(iterator); + ASSERT_TRUE(iterator->is_valid()); + const auto *stored = static_cast(iterator->data()); + for (size_t i = 0; i < kDimension; ++i) { + EXPECT_EQ(static_cast(vectors[0][i]), + static_cast(static_cast(stored[i]) + 128)); + } + EXPECT_EQ(SumSquared(vectors[0]), ReadTail(stored)); + + auto reformer = IndexFactory::CreateReformer("UniformUint8StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(converter->meta().reformer_params())); + + std::string record; + IndexQueryMeta record_meta; + ASSERT_EQ(0, reformer->convert( + vectors[0].data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension), + &record, &record_meta)); + EXPECT_EQ(std::string(reinterpret_cast(stored), + converted_holder->element_size()), + record); + + std::string reverted; + ASSERT_EQ(0, reformer->revert(record.data(), record_meta, &reverted)); + ASSERT_EQ(kDimension * sizeof(float), reverted.size()); + for (size_t i = 0; i < kDimension; ++i) { + EXPECT_FLOAT_EQ(vectors[0][i], ReadFloat(reverted, i)); + } +} + +TEST(UniformUint8Reformer, UsesAsymmetricRecordAndQueryLayouts) { + const std::vector> vectors{ + {0.0f, 1.0f, 127.0f, 128.0f, 254.0f, 255.0f}, + {255.0f, 0.0f, 200.0f, 100.0f, 50.0f, 25.0f}, + }; + auto holder = MakeHolder(vectors); + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto converter = + IndexFactory::CreateConverter("UniformUint8StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0, converter->init(meta, ailego::Params())); + ASSERT_EQ(0, converter->train(holder)); + + auto reformer = IndexFactory::CreateReformer("UniformUint8StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(converter->meta().reformer_params())); + + std::string record; + std::string query; + IndexQueryMeta record_meta; + IndexQueryMeta query_meta; + const IndexQueryMeta input_meta(IndexMeta::DataType::DT_FP32, kDimension); + ASSERT_EQ(0, reformer->convert(vectors[0].data(), input_meta, &record, + &record_meta)); + ASSERT_EQ(0, reformer->transform(vectors[0].data(), input_meta, &query, + &query_meta)); + EXPECT_EQ(kEncodedDimension, record_meta.dimension()); + EXPECT_EQ(kEncodedDimension, query_meta.dimension()); + EXPECT_EQ(record_meta.element_size(), query_meta.element_size()); + EXPECT_NE(record, query); + + const auto *query_bytes = reinterpret_cast(query.data()); + for (size_t i = 0; i < kDimension; ++i) { + const auto record_code = static_cast( + static_cast(static_cast(record[i])) + 128); + EXPECT_EQ(static_cast(vectors[0][i]), record_code); + EXPECT_EQ(static_cast(vectors[0][i]), query_bytes[i]); + } + EXPECT_EQ(SumSquared(vectors[0]), ReadTail(record.data())); + EXPECT_EQ(SumSquared(vectors[0]), ReadTail(query.data())); +} + +TEST(UniformUint8Reformer, BatchEncodingUsesEncodedStride) { + const std::vector> vectors{ + {0.0f, 1.0f, 127.0f, 128.0f, 254.0f, 255.0f}, + {255.0f, 0.0f, 200.0f, 100.0f, 50.0f, 25.0f}, + }; + auto holder = MakeHolder(vectors); + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto converter = + IndexFactory::CreateConverter("UniformUint8StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0, converter->init(meta, ailego::Params())); + ASSERT_EQ(0, converter->train(holder)); + + auto reformer = IndexFactory::CreateReformer("UniformUint8StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(converter->meta().reformer_params())); + + const std::vector inputs{ + 0.0f, 1.0f, 127.0f, 128.0f, 254.0f, 255.0f, + 255.0f, 0.0f, 200.0f, 100.0f, 50.0f, 25.0f, + }; + const IndexQueryMeta input_meta(IndexMeta::DataType::DT_FP32, kDimension); + std::string records; + std::string queries; + IndexQueryMeta record_meta; + IndexQueryMeta query_meta; + ASSERT_EQ(0, reformer->convert(inputs.data(), input_meta, 2, &records, + &record_meta)); + ASSERT_EQ(0, reformer->transform(inputs.data(), input_meta, 2, &queries, + &query_meta)); + ASSERT_EQ(2 * kEncodedDimension, records.size()); + ASSERT_EQ(2 * kEncodedDimension, queries.size()); + + for (size_t vector_index = 0; vector_index < vectors.size(); ++vector_index) { + const size_t offset = vector_index * kEncodedDimension; + const auto *record = + reinterpret_cast(records.data() + offset); + const auto *query = + reinterpret_cast(queries.data() + offset); + for (size_t i = 0; i < kDimension; ++i) { + EXPECT_EQ(static_cast(vectors[vector_index][i]), + static_cast(static_cast(record[i]) + 128)); + EXPECT_EQ(static_cast(vectors[vector_index][i]), query[i]); + } + EXPECT_EQ(SumSquared(vectors[vector_index]), ReadTail(record)); + EXPECT_EQ(SumSquared(vectors[vector_index]), ReadTail(query)); + } +} + +TEST(UniformUint8Reformer, RestoresParamsFromConverterMetaWithoutRetraining) { + const std::vector> vectors{ + {-4.5f, -2.0f, 0.5f, 3.0f, 5.5f, 8.0f}, + {8.0f, 5.5f, 3.0f, 0.5f, -2.0f, -4.5f}, + }; + auto holder = MakeHolder(vectors); + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto trained = + IndexFactory::CreateConverter("UniformUint8StreamingConverter"); + ASSERT_TRUE(trained); + ASSERT_EQ(0, trained->init(meta, ailego::Params())); + ASSERT_EQ(0, trained->train(holder)); + + auto restored = + IndexFactory::CreateConverter("UniformUint8StreamingConverter"); + ASSERT_TRUE(restored); + ASSERT_EQ(0, restored->init(meta, trained->meta().converter_params())); + EXPECT_EQ("UniformUint8StreamingReformer", restored->meta().reformer_name()); + ASSERT_EQ(0, restored->transform(holder)); + EXPECT_EQ(vectors.size(), restored->stats().transformed_count()); + ASSERT_TRUE(restored->result()); + EXPECT_EQ(kEncodedDimension, restored->result()->dimension()); + + ASSERT_EQ(0, restored->cleanup()); + ASSERT_EQ(0, restored->init(meta, ailego::Params())); + EXPECT_NE(0, restored->transform(holder)); +} + +TEST(UniformUint8Reformer, RejectsEmptyParamsAndResetsInitializedState) { + ailego::Params params; + params.set("uniform_uint8.reformer.scale", 1.0f); + params.set("uniform_uint8.reformer.bias", 0.0f); + + auto reformer = IndexFactory::CreateReformer("UniformUint8StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(params)); + + const std::vector input(kDimension, 1.0f); + const IndexQueryMeta input_meta(IndexMeta::DataType::DT_FP32, kDimension); + std::string output; + IndexQueryMeta output_meta; + ASSERT_EQ( + 0, reformer->transform(input.data(), input_meta, &output, &output_meta)); + + ASSERT_EQ(IndexError_InvalidArgument, reformer->init(ailego::Params())); + EXPECT_NE( + 0, reformer->transform(input.data(), input_meta, &output, &output_meta)); +} + +TEST(UniformUint8Reformer, RejectsNonFiniteTrainingData) { + const std::vector> vectors{ + {0.0f, 1.0f, 2.0f, std::numeric_limits::infinity(), 4.0f, 5.0f}, + }; + auto holder = MakeHolder(vectors); + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto converter = + IndexFactory::CreateConverter("UniformUint8StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0, converter->init(meta, ailego::Params())); + EXPECT_NE(0, converter->train(holder)); +} + +TEST(UniformUint8Reformer, StoresUint32SquaredSumAtMaximumDimension) { + constexpr size_t kMaxDimension = MAX_DIMENSION; + std::vector input(kMaxDimension, 255.0f); + + ailego::Params params; + params.set("uniform_uint8.reformer.scale", 1.0f); + params.set("uniform_uint8.reformer.bias", 0.0f); + auto reformer = IndexFactory::CreateReformer("UniformUint8StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(params)); + + std::string output; + IndexQueryMeta output_meta; + ASSERT_EQ(0, reformer->convert( + input.data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kMaxDimension), + &output, &output_meta)); + EXPECT_EQ(kMaxDimension + sizeof(uint32_t), output_meta.dimension()); + EXPECT_EQ(uint64_t{kMaxDimension} * 255 * 255, + ReadTail(output.data(), kMaxDimension)); +} + +TEST(UniformUint8Reformer, RejectsDimensionsOutsidePublicRange) { + ailego::Params params; + params.set("uniform_uint8.reformer.scale", 1.0f); + params.set("uniform_uint8.reformer.bias", 0.0f); + auto reformer = IndexFactory::CreateReformer("UniformUint8StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(params)); + + const float input = 0.0f; + std::string output; + IndexQueryMeta output_meta; + for (const uint32_t dimension : + {0U, static_cast(MAX_DIMENSION + 1)}) { + EXPECT_NE( + 0, reformer->convert( + &input, IndexQueryMeta(IndexMeta::DataType::DT_FP32, dimension), + &output, &output_meta)) + << "dimension=" << dimension; + } + + IndexMeta oversized_meta(IndexMeta::DataType::DT_FP32, MAX_DIMENSION + 1); + oversized_meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto converter = + IndexFactory::CreateConverter("UniformUint8StreamingConverter"); + ASSERT_TRUE(converter); + EXPECT_NE(0, converter->init(oversized_meta, ailego::Params())); +} + +} // namespace +} // namespace zvec::core From 3baed4a753e35ec689e75add1be11a5f95c5cb08 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Thu, 30 Jul 2026 20:20:02 +0800 Subject: [PATCH 2/2] fix(quantization): validate uint8 revert metadata --- src/core/quantizer/uniform_uint8_reformer.cc | 10 ++++++- .../quantizer/uniform_uint8_reformer_test.cc | 29 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/core/quantizer/uniform_uint8_reformer.cc b/src/core/quantizer/uniform_uint8_reformer.cc index 71ca180e1..d385dbeeb 100644 --- a/src/core/quantizer/uniform_uint8_reformer.cc +++ b/src/core/quantizer/uniform_uint8_reformer.cc @@ -106,7 +106,9 @@ class UniformUint8StreamingReformer : public IndexReformer { int revert(const void *input, const IndexQueryMeta &record_meta, std::string *output) const override { if (!initialized_ || !input || !output || - record_meta.data_type() != IndexMeta::DataType::DT_INT8) { + record_meta.data_type() != IndexMeta::DataType::DT_INT8 || + record_meta.unit_size() != + IndexMeta::UnitSizeof(IndexMeta::DataType::DT_INT8)) { return IndexError_InvalidArgument; } @@ -114,6 +116,12 @@ class UniformUint8StreamingReformer : public IndexReformer { return IndexError_InvalidArgument; } const size_t dimension = record_meta.dimension() - kTailBytes; + if (dimension > MAX_DIMENSION) { + LOG_ERROR( + "UniformUint8StreamingReformer: dimension=%zu must be in [1, %d]", + dimension, MAX_DIMENSION); + return IndexError_InvalidArgument; + } output->resize(dimension * sizeof(float)); auto *destination = output->data(); const auto *source = static_cast(input); diff --git a/tests/core/quantizer/uniform_uint8_reformer_test.cc b/tests/core/quantizer/uniform_uint8_reformer_test.cc index da0a62d3d..0e15183ac 100644 --- a/tests/core/quantizer/uniform_uint8_reformer_test.cc +++ b/tests/core/quantizer/uniform_uint8_reformer_test.cc @@ -338,5 +338,34 @@ TEST(UniformUint8Reformer, RejectsDimensionsOutsidePublicRange) { EXPECT_NE(0, converter->init(oversized_meta, ailego::Params())); } +TEST(UniformUint8Reformer, RejectsMalformedRevertMetadata) { + ailego::Params params; + params.set("uniform_uint8.reformer.scale", 1.0f); + params.set("uniform_uint8.reformer.bias", 0.0f); + auto reformer = IndexFactory::CreateReformer("UniformUint8StreamingReformer"); + ASSERT_TRUE(reformer); + ASSERT_EQ(0, reformer->init(params)); + + const std::vector record(kEncodedDimension, 0); + std::string output; + + const IndexQueryMeta invalid_unit_size( + IndexMeta::MetaType::MT_DENSE, IndexMeta::DataType::DT_INT8, + IndexMeta::UnitSizeof(IndexMeta::DataType::DT_INT8) + 1, + kEncodedDimension); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->revert(record.data(), invalid_unit_size, &output)); + + const IndexQueryMeta missing_vector_data(IndexMeta::DataType::DT_INT8, + sizeof(uint32_t)); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->revert(record.data(), missing_vector_data, &output)); + + const IndexQueryMeta oversized_dimension( + IndexMeta::DataType::DT_INT8, MAX_DIMENSION + sizeof(uint32_t) + 1); + EXPECT_EQ(IndexError_InvalidArgument, + reformer->revert(record.data(), oversized_dimension, &output)); +} + } // namespace } // namespace zvec::core