From ff2010fb8be37f7629fd0bc0599bd192f16e5cd4 Mon Sep 17 00:00:00 2001 From: Jianning Wang Date: Tue, 25 Aug 2026 14:44:52 +0800 Subject: [PATCH 1/3] feat(turbo): init uniform uint4 --- src/core/interface/index.cc | 3 + src/core/metric/metric_params.h | 4 + src/core/metric/uniform_uint4_metric.cc | 132 ++++++ src/core/quantizer/quantizer_params.h | 8 + src/core/quantizer/uniform_uint4_converter.cc | 415 ++++++++++++++++++ src/core/quantizer/uniform_uint4_reformer.cc | 212 +++++++++ src/include/zvec/core/interface/index_param.h | 2 + src/include/zvec/turbo/turbo.h | 11 + .../avx512_vnni/uniform_uint4/quantize.cc | 50 +++ .../avx512_vnni/uniform_uint4/quantize.h | 13 + .../uniform_uint4/squared_euclidean.cc | 109 +++++ .../uniform_uint4/squared_euclidean.h | 20 + src/turbo/turbo.cc | 16 + tests/core/interface/index_interface_test.cc | 2 + .../core/metric/uniform_uint4_metric_test.cc | 102 +++++ .../quantizer/uniform_uint4_reformer_test.cc | 154 +++++++ 16 files changed, 1253 insertions(+) create mode 100644 src/core/metric/uniform_uint4_metric.cc create mode 100644 src/core/quantizer/uniform_uint4_converter.cc create mode 100644 src/core/quantizer/uniform_uint4_reformer.cc create mode 100644 src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc create mode 100644 src/turbo/distance/avx512_vnni/uniform_uint4/quantize.h create mode 100644 src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc create mode 100644 src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.h create mode 100644 tests/core/metric/uniform_uint4_metric_test.cc create mode 100644 tests/core/quantizer/uniform_uint4_reformer_test.cc diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 557871fc0..d0c283333 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -218,6 +218,9 @@ int Index::CreateAndInitConverterReformer(const QuantizerParam ¶m, case QuantizerType::kUniformUint8: converter_name = "UniformUint8Converter"; break; + case QuantizerType::kUniformUint4: + converter_name = "UniformUint4Converter"; + break; default: LOG_ERROR("Unsupported quantizer type: "); return core::IndexError_Unsupported; diff --git a/src/core/metric/metric_params.h b/src/core/metric/metric_params.h index 824e31a33..2aa26e5aa 100644 --- a/src/core/metric/metric_params.h +++ b/src/core/metric/metric_params.h @@ -42,5 +42,9 @@ static const std::string UNIFORM_UINT7_METRIC_ORIGIN_METRIC_NAME = static const std::string UNIFORM_UINT8_METRIC_ORIGIN_METRIC_NAME = "proxima.uniform_uint8.metric.origin_metric_name"; +//! UniformUint4 Metric +static const std::string UNIFORM_UINT4_METRIC_ORIGIN_METRIC_NAME = + "proxima.uniform_uint4.metric.origin_metric_name"; + } // namespace core } // namespace zvec diff --git a/src/core/metric/uniform_uint4_metric.cc b/src/core/metric/uniform_uint4_metric.cc new file mode 100644 index 000000000..4275130b2 --- /dev/null +++ b/src/core/metric/uniform_uint4_metric.cc @@ -0,0 +1,132 @@ +// 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 "metric_params.h" + +namespace zvec { +namespace core { +namespace { + +void UniformUint4SquaredEuclidean(const void *lhs, const void *rhs, + size_t encoded_dimension, float *distance) { + const auto *a = static_cast(lhs); + const auto *b = static_cast(rhs); + int64_t sum = 0; + for (size_t i = 0; i < encoded_dimension; ++i) { + const int low_delta = + static_cast(a[i] & 0x0fU) - static_cast(b[i] & 0x0fU); + const int high_delta = static_cast((a[i] >> 4U) & 0x0fU) - + static_cast((b[i] >> 4U) & 0x0fU); + sum += low_delta * low_delta + high_delta * high_delta; + } + *distance = static_cast(sum); +} + +void UniformUint4SquaredEuclideanBatch(const void *const *vectors, + const void *query, size_t count, + size_t encoded_dimension, + float *distances) { + for (size_t i = 0; i < count; ++i) { + UniformUint4SquaredEuclidean(vectors[i], query, encoded_dimension, + distances + i); + } +} + +} // namespace + +class UniformUint4Metric : public IndexMetric { + public: + int init(const IndexMeta &meta, const ailego::Params ¶ms) override { + if (meta.data_type() != IndexMeta::DataType::DT_INT8 || + meta.dimension() == 0 || (meta.dimension() % 64U) != 0) { + LOG_ERROR( + "UniformUint4Metric: expected a non-empty packed DT_INT8 dimension " + "aligned to 64 bytes, got type=%d dimension=%u", + meta.data_type(), meta.dimension()); + return IndexError_Unsupported; + } + std::string origin_metric; + params.get(UNIFORM_UINT4_METRIC_ORIGIN_METRIC_NAME, &origin_metric); + if (origin_metric != "SquaredEuclidean") { + LOG_ERROR("UniformUint4Metric: only SquaredEuclidean is supported"); + 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.dimension() == meta_.dimension() && + meta.unit_size() == meta_.unit_size(); + } + bool is_matched(const IndexMeta &meta, + const IndexQueryMeta &qmeta) const override { + return is_matched(meta) && qmeta.data_type() == meta_.data_type() && + qmeta.dimension() == meta_.dimension() && + qmeta.unit_size() == meta_.unit_size(); + } + + MatrixDistance distance(void) const override { + auto turbo_distance = turbo::get_distance_func( + turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt4, + turbo::QuantizeType::kUniformUint4); + return turbo_distance ? turbo_distance : UniformUint4SquaredEuclidean; + } + MatrixDistance distance_matrix(size_t m, size_t n) const override { + return m == 1 && n == 1 ? distance() : MatrixDistance{}; + } + MatrixBatchDistance batch_distance(void) const override { + auto turbo_distance = turbo::get_batch_distance_func( + turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt4, + turbo::QuantizeType::kUniformUint4); + return turbo_distance ? turbo_distance : UniformUint4SquaredEuclideanBatch; + } + DistanceBatchQueryPreprocessFunc get_query_preprocess_func() const override { + return nullptr; + } + 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; + } + + private: + IndexMeta meta_{}; + ailego::Params params_{}; +}; + +INDEX_FACTORY_REGISTER_METRIC_ALIAS(UniformUint4, UniformUint4Metric); + +} // namespace core +} // namespace zvec diff --git a/src/core/quantizer/quantizer_params.h b/src/core/quantizer/quantizer_params.h index 5d706954f..ab30ae2bf 100644 --- a/src/core/quantizer/quantizer_params.h +++ b/src/core/quantizer/quantizer_params.h @@ -131,6 +131,14 @@ static const std::string UNIFORM_UINT8_REFORMER_SCALE = static const std::string UNIFORM_UINT8_REFORMER_BIAS = "uniform_uint8.reformer.bias"; +//! UniformUint4Converter / Reformer +static const std::string UNIFORM_UINT4_REFORMER_MINIMUM = + "uniform_uint4.reformer.minimum"; +static const std::string UNIFORM_UINT4_REFORMER_RANGE = + "uniform_uint4.reformer.range"; +static const std::string UNIFORM_UINT4_REFORMER_ORIGINAL_DIMENSION = + "uniform_uint4.reformer.original_dimension"; + //! DoubleBitConverter static const std::string DOUBLE_BIT_CONVERTER_TRAIN_SAMPLE_COUNT = "double_bit.converter.train_sample_count"; diff --git a/src/core/quantizer/uniform_uint4_converter.cc b/src/core/quantizer/uniform_uint4_converter.cc new file mode 100644 index 000000000..57b8eb417 --- /dev/null +++ b/src/core/quantizer/uniform_uint4_converter.cc @@ -0,0 +1,415 @@ +// 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 +#include "../metric/metric_params.h" + +namespace zvec { +namespace core { +namespace { + +constexpr float kAlmostHalf = 0.4999999701976776123046875f; + +size_t PaddedDimension(size_t dimension) { + return (dimension + 127U) / 128U * 128U; +} + +uint32_t FloatOrderKey(float value) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return (bits & 0x80000000U) != 0 ? ~bits : bits ^ 0x80000000U; +} + +float FloatFromOrderKey(uint32_t key) { + const uint32_t bits = (key & 0x80000000U) != 0 ? key ^ 0x80000000U : ~key; + float value = 0.0f; + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +struct RadixSelection { + size_t rank{0}; + uint32_t prefix{0}; +}; + +bool ChooseBucket(int shift, const std::array &histogram, + RadixSelection *selection) { + size_t before = 0; + for (size_t bucket = 0; bucket < histogram.size(); ++bucket) { + if (selection->rank < before + histogram[bucket]) { + selection->rank -= before; + selection->prefix |= static_cast(bucket) << shift; + return true; + } + before += histogram[bucket]; + } + return false; +} + +void QuantizeScalar(const float *input, size_t dimension, float minimum, + float range, uint8_t *output, size_t encoded_dimension) { + std::memset(output, 0, encoded_dimension); + for (size_t d = 0; d < dimension; ++d) { + float normalized = (input[d] - minimum) / range; + normalized = std::min(1.0f, std::max(0.0f, normalized)); + const auto code = static_cast( + static_cast(normalized * 15.0f + kAlmostHalf)); + if ((d & 1U) == 0) { + output[d >> 1U] = code; + } else { + output[d >> 1U] |= static_cast(code << 4U); + } + } +} + +bool IsSupportedSourceType(IndexMeta::DataType data_type) { + return data_type == IndexMeta::DataType::DT_FP32 || + data_type == IndexMeta::DataType::DT_FP16; +} + +float SourceValue(const void *record, IndexMeta::DataType data_type, + size_t index) { + switch (data_type) { + case IndexMeta::DataType::DT_FP32: + return static_cast(record)[index]; + case IndexMeta::DataType::DT_FP16: + return static_cast( + static_cast(record)[index]); + default: + return std::numeric_limits::quiet_NaN(); + } +} + +void DecodeSource(const void *record, IndexMeta::DataType data_type, + size_t dimension, std::vector *output) { + output->resize(dimension); + for (size_t i = 0; i < dimension; ++i) { + (*output)[i] = SourceValue(record, data_type, i); + } +} + +} // namespace + +class UniformUint4Converter : public IndexConverter { + public: + UniformUint4Converter(IndexMeta::DataType /*dst_type*/) {} + ~UniformUint4Converter() override = default; + + int init(const IndexMeta &index_meta, const ailego::Params ¶ms) override { + if (index_meta.data_type() != IndexMeta::DataType::DT_FP32 || + index_meta.dimension() == 0) { + return IndexError_Unsupported; + } + meta_ = index_meta; + original_dimension_ = index_meta.dimension(); + encoded_dimension_ = PaddedDimension(original_dimension_) / 2U; + *stats_.mutable_trained_count() = 0; + *stats_.mutable_transformed_count() = 0; + + meta_.set_converter("UniformUint4Converter", 0, params); + meta_.set_meta(IndexMeta::DataType::DT_INT8, encoded_dimension_); + + ailego::Params metric_params; + metric_params.set(UNIFORM_UINT4_METRIC_ORIGIN_METRIC_NAME, + index_meta.metric_name()); + meta_.set_metric("UniformUint4", 0, metric_params); + + const bool has_minimum = + params.get(UNIFORM_UINT4_REFORMER_MINIMUM, &minimum_); + const bool has_range = params.get(UNIFORM_UINT4_REFORMER_RANGE, &range_); + if (has_minimum && has_range && range_ > 0.0f && std::isfinite(minimum_) && + std::isfinite(range_)) { + SetReformerParams(); + } + return 0; + } + + int cleanup(void) override { + *stats_.mutable_trained_count() = 0; + *stats_.mutable_transformed_count() = 0; + holder_.reset(); + return 0; + } + + int train(IndexHolder::Pointer holder) override { + if (!holder || !IsSupportedSourceType(holder->data_type()) || + holder->dimension() != original_dimension_) { + return IndexError_Mismatch; + } + const auto source_type = holder->data_type(); + + ailego::ElapsedTime timer; + AILEGO_DEFER([&]() { stats_.set_trained_costtime(timer.milli_seconds()); }); + + size_t record_count = holder->count(); + if (record_count == 0) { + LOG_ERROR("UniformUint4Converter: empty training set"); + return IndexError_InvalidArgument; + } + if (record_count == std::numeric_limits::max() || + record_count > + std::numeric_limits::max() / original_dimension_) { + LOG_ERROR("UniformUint4Converter: invalid training count"); + return IndexError_InvalidArgument; + } + + if (holder->multipass()) { + const size_t value_count = record_count * original_dimension_; + // Match reimpl/vamana and KGN exactly: discard + // floor(float(N) * 0.01) + 1 values at each tail. + size_t tail = + static_cast(static_cast(value_count) * 0.01f) + + size_t{1}; + tail = std::max(1, std::min(tail, value_count)); + RadixSelection lower{tail - 1, 0}; + RadixSelection upper{value_count - tail, 0}; + + for (int shift = 24; shift >= 0; shift -= 8) { + std::array lower_hist{}; + std::array upper_hist{}; + const uint32_t prefix_mask = + shift == 24 ? 0U : ~uint32_t{0} << (shift + 8); + auto iter = holder->create_iterator(); + if (!iter) { + LOG_ERROR("UniformUint4Converter: iterator unavailable"); + return IndexError_Runtime; + } + size_t actual_records = 0; + for (; iter->is_valid(); iter->next(), ++actual_records) { + for (size_t d = 0; d < original_dimension_; ++d) { + const float value = SourceValue(iter->data(), source_type, d); + if (!std::isfinite(value)) { + LOG_ERROR( + "UniformUint4Converter: non-finite training " + "value (record_idx=%zu, dim_idx=%zu)", + actual_records, d); + return IndexError_InvalidArgument; + } + const uint32_t key = FloatOrderKey(value); + const size_t bucket = (key >> shift) & 0xffU; + if ((key & prefix_mask) == lower.prefix) ++lower_hist[bucket]; + if ((key & prefix_mask) == upper.prefix) ++upper_hist[bucket]; + } + } + if (actual_records != record_count || + !ChooseBucket(shift, lower_hist, &lower) || + !ChooseBucket(shift, upper_hist, &upper)) { + LOG_ERROR("UniformUint4Converter: radix selection failed"); + return IndexError_Runtime; + } + } + minimum_ = FloatFromOrderKey(lower.prefix); + const float maximum = FloatFromOrderKey(upper.prefix); + range_ = maximum - minimum_; + } else { + // IndexConverter wraps one-shot sources in a two-pass holder. Such a + // holder cannot support exact four-pass order-statistic selection, so + // retain bounded-memory compatibility by using its exact global range. + LOG_WARN( + "UniformUint4Converter: one-pass holder; using global " + "min/max instead of 1%% clipped calibration"); + float maximum = std::numeric_limits::lowest(); + minimum_ = std::numeric_limits::max(); + auto iter = holder->create_iterator(); + if (!iter) return IndexError_Runtime; + record_count = 0; + for (; iter->is_valid(); iter->next(), ++record_count) { + for (size_t d = 0; d < original_dimension_; ++d) { + const float value = SourceValue(iter->data(), source_type, d); + if (!std::isfinite(value)) return IndexError_InvalidArgument; + minimum_ = std::min(minimum_, value); + maximum = std::max(maximum, value); + } + } + range_ = maximum - minimum_; + } + + if (!(range_ > 0.0f)) range_ = 1.0f; + *stats_.mutable_trained_count() = record_count; + SetReformerParams(); + + ailego::Params converter_params = meta_.converter_params(); + converter_params.set(UNIFORM_UINT4_REFORMER_MINIMUM, minimum_); + converter_params.set(UNIFORM_UINT4_REFORMER_RANGE, range_); + converter_params.set(UNIFORM_UINT4_REFORMER_ORIGINAL_DIMENSION, + original_dimension_); + meta_.set_converter(meta_.converter_name(), 0, converter_params); + LOG_INFO( + "UniformUint4Converter train done: costtime %zums, " + "minimum=%f, range=%f, original_dimension=%zu, encoded_dimension=%zu", + static_cast(timer.milli_seconds()), minimum_, range_, + original_dimension_, encoded_dimension_); + return 0; + } + + int transform(IndexHolder::Pointer holder) override { + if (!holder || !IsSupportedSourceType(holder->data_type()) || + holder->dimension() != original_dimension_ || !(range_ > 0.0f)) { + return IndexError_Mismatch; + } + if (holder->count() > 0) { + *stats_.mutable_transformed_count() += holder->count(); + } + holder_ = std::make_shared( + std::move(holder), original_dimension_, encoded_dimension_, minimum_, + range_); + 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: + void SetReformerParams() { + ailego::Params reformer_params; + reformer_params.set(UNIFORM_UINT4_REFORMER_MINIMUM, minimum_); + reformer_params.set(UNIFORM_UINT4_REFORMER_RANGE, range_); + reformer_params.set(UNIFORM_UINT4_REFORMER_ORIGINAL_DIMENSION, + original_dimension_); + meta_.set_reformer("UniformUint4Reformer", 0, reformer_params); + } + + class UniformUint4Holder : public IndexHolder { + public: + class Iterator : public IndexHolder::Iterator { + public: + Iterator(const UniformUint4Holder *owner, + IndexHolder::Iterator::Pointer &&front) + : owner_(owner), + buffer_(owner->encoded_dimension_, 0), + front_(std::move(front)) { + Encode(); + } + + const void *data(void) const override { + return buffer_.data(); + } + bool is_valid(void) const override { + return front_->is_valid(); + } + uint64_t key(void) const override { + return front_->key(); + } + void next(void) override { + front_->next(); + Encode(); + } + + private: + void Encode() { + if (!front_->is_valid()) return; + const float *input = nullptr; + if (owner_->source_type_ == IndexMeta::DataType::DT_FP32) { + input = static_cast(front_->data()); + } else { + DecodeSource(front_->data(), owner_->source_type_, + owner_->original_dimension_, &decoded_); + input = decoded_.data(); + } + if (owner_->quantize_func_) { + owner_->quantize_func_(input, owner_->original_dimension_, + owner_->minimum_, owner_->range_, + buffer_.data()); + } else { + QuantizeScalar(input, owner_->original_dimension_, owner_->minimum_, + owner_->range_, buffer_.data(), + owner_->encoded_dimension_); + } + } + + const UniformUint4Holder *owner_{nullptr}; + std::vector buffer_{}; + std::vector decoded_{}; + IndexHolder::Iterator::Pointer front_{}; + }; + + UniformUint4Holder(IndexHolder::Pointer front, size_t original_dimension, + size_t encoded_dimension, float minimum, float range) + : front_(std::move(front)), + source_type_(front_->data_type()), + original_dimension_(original_dimension), + encoded_dimension_(encoded_dimension), + minimum_(minimum), + range_(range), + quantize_func_( + turbo::get_uniform_uint4_quantize_func(turbo::DataType::kInt4)) {} + + size_t count(void) const override { + return front_->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 encoded_dimension_; + } + bool multipass(void) const override { + return front_->multipass(); + } + IndexHolder::Iterator::Pointer create_iterator(void) override { + auto iter = front_->create_iterator(); + return iter ? IndexHolder::Iterator::Pointer( + new Iterator(this, std::move(iter))) + : IndexHolder::Iterator::Pointer(); + } + + private: + IndexHolder::Pointer front_{}; + IndexMeta::DataType source_type_{IndexMeta::DataType::DT_UNDEFINED}; + size_t original_dimension_{0}; + size_t encoded_dimension_{0}; + float minimum_{0.0f}; + float range_{0.0f}; + turbo::UniformUint4QuantizeFunc quantize_func_{nullptr}; + }; + + IndexMeta meta_{}; + Stats stats_{}; + IndexHolder::Pointer holder_{}; + size_t original_dimension_{0}; + size_t encoded_dimension_{0}; + float minimum_{0.0f}; + float range_{0.0f}; +}; + +INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(UniformUint4Converter, + UniformUint4Converter, + IndexMeta::DataType::DT_INT8); + +} // namespace core +} // namespace zvec diff --git a/src/core/quantizer/uniform_uint4_reformer.cc b/src/core/quantizer/uniform_uint4_reformer.cc new file mode 100644 index 000000000..d6fc9f8c0 --- /dev/null +++ b/src/core/quantizer/uniform_uint4_reformer.cc @@ -0,0 +1,212 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace core { +namespace { + +constexpr float kAlmostHalf = 0.4999999701976776123046875f; + +size_t EncodedDimension(size_t dimension) { + return ((dimension + 127U) / 128U * 128U) / 2U; +} + +void QuantizeScalar(const float *input, size_t dimension, float minimum, + float range, uint8_t *output, size_t encoded_dimension) { + std::memset(output, 0, encoded_dimension); + for (size_t d = 0; d < dimension; ++d) { + float normalized = (input[d] - minimum) / range; + normalized = std::min(1.0f, std::max(0.0f, normalized)); + const auto code = static_cast( + static_cast(normalized * 15.0f + kAlmostHalf)); + if ((d & 1U) == 0) { + output[d >> 1U] = code; + } else { + output[d >> 1U] |= static_cast(code << 4U); + } + } +} + +} // namespace + +class UniformUint4Reformer : public IndexReformer { + public: + UniformUint4Reformer(IndexMeta::DataType /*dst_type*/) {} + ~UniformUint4Reformer() override = default; + + int init(const ailego::Params ¶ms) override { + uint32_t original_dimension = 0; + const bool has_minimum = + params.get(UNIFORM_UINT4_REFORMER_MINIMUM, &minimum_); + const bool has_range = params.get(UNIFORM_UINT4_REFORMER_RANGE, &range_); + const bool has_dimension = params.get( + UNIFORM_UINT4_REFORMER_ORIGINAL_DIMENSION, &original_dimension); + if (!has_minimum || !has_range || !has_dimension || + !std::isfinite(minimum_) || !std::isfinite(range_) || + !(range_ > 0.0f) || original_dimension == 0) { + LOG_ERROR("UniformUint4Reformer: invalid or missing params"); + return IndexError_InvalidArgument; + } + original_dimension_ = original_dimension; + encoded_dimension_ = EncodedDimension(original_dimension_); + const float step = range_ / 15.0f; + distance_scale_ = step * step; + quantize_func_ = + turbo::get_uniform_uint4_quantize_func(turbo::DataType::kInt4); + initialized_ = true; + return 0; + } + + int cleanup(void) override { + initialized_ = false; + return 0; + } + int load(IndexStorage::Pointer) override { + return 0; + } + int unload(void) override { + return 0; + } + + int transform(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const override { + return Quantize(query, qmeta, 1, false, out, ometa); + } + int transform(const void *query, const IndexQueryMeta &qmeta, uint32_t count, + std::string *out, IndexQueryMeta *ometa) const override { + return Quantize(query, qmeta, count, false, out, ometa); + } + int convert(const void *record, const IndexQueryMeta &rmeta, std::string *out, + IndexQueryMeta *ometa) const override { + return Quantize(record, rmeta, 1, true, out, ometa); + } + int convert(const void *records, const IndexQueryMeta &rmeta, uint32_t count, + std::string *out, IndexQueryMeta *ometa) const override { + return Quantize(records, rmeta, count, true, out, ometa); + } + + int normalize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, + IndexDocumentList &result) const override { + if (!initialized_) return IndexError_Runtime; + for (auto &item : result) { + *item.mutable_score() *= distance_scale_; + } + return 0; + } + + bool need_revert() const override { + return true; + } + + int revert(const void *input, const IndexQueryMeta &qmeta, + std::string *out) const override { + if (!initialized_) return IndexError_Runtime; + if (qmeta.data_type() != IndexMeta::DataType::DT_INT8 || + qmeta.dimension() != encoded_dimension_) { + return IndexError_Mismatch; + } + out->resize(original_dimension_ * sizeof(float)); + auto *decoded = reinterpret_cast(out->data()); + const auto *packed = static_cast(input); + const float step = range_ / 15.0f; + for (size_t d = 0; d < original_dimension_; ++d) { + const uint8_t byte = packed[d >> 1U]; + const uint8_t code = (d & 1U) == 0 ? byte & 0x0fU : (byte >> 4U) & 0x0fU; + decoded[d] = minimum_ + static_cast(code) * step; + } + return 0; + } + + private: + int Quantize(const void *source, const IndexQueryMeta &source_meta, + uint32_t count, bool accept_native_flat, std::string *out, + IndexQueryMeta *output_meta) const { + if (!initialized_) return IndexError_Runtime; + const auto source_type = source_meta.data_type(); + const bool is_fp32 = source_type == IndexMeta::DataType::DT_FP32; + const bool is_native_flat = + accept_native_flat && source_type == IndexMeta::DataType::DT_FP16; + if (!source || !out || !output_meta || (!is_fp32 && !is_native_flat) || + source_meta.dimension() != original_dimension_) { + return IndexError_Mismatch; + } + + *output_meta = source_meta; + output_meta->set_meta(IndexMeta::DataType::DT_INT8, encoded_dimension_); + const size_t output_stride = output_meta->element_size(); + out->resize(static_cast(count) * output_stride); + auto *output = reinterpret_cast(out->data()); + const auto *source_bytes = static_cast(source); + static thread_local std::vector decoded; + if (!is_fp32) decoded.resize(original_dimension_); + for (uint32_t i = 0; i < count; ++i) { + const void *source_row = + source_bytes + static_cast(i) * source_meta.element_size(); + const float *row = nullptr; + if (is_fp32) { + row = static_cast(source_row); + } else if (source_type == IndexMeta::DataType::DT_FP16) { + const auto *input = static_cast(source_row); + for (size_t d = 0; d < original_dimension_; ++d) { + decoded[d] = static_cast(input[d]); + } + row = decoded.data(); + } else { + const auto *input = static_cast(source_row); + for (size_t d = 0; d < original_dimension_; ++d) { + decoded[d] = static_cast(input[d]); + } + row = decoded.data(); + } + for (size_t d = 0; d < original_dimension_; ++d) { + if (!std::isfinite(row[d])) { + LOG_ERROR("UniformUint4Reformer: non-finite input value"); + return IndexError_InvalidArgument; + } + } + uint8_t *encoded = output + static_cast(i) * output_stride; + if (quantize_func_) { + quantize_func_(row, original_dimension_, minimum_, range_, encoded); + } else { + QuantizeScalar(row, original_dimension_, minimum_, range_, encoded, + encoded_dimension_); + } + } + return 0; + } + + float minimum_{0.0f}; + float range_{0.0f}; + float distance_scale_{1.0f}; + size_t original_dimension_{0}; + size_t encoded_dimension_{0}; + bool initialized_{false}; + turbo::UniformUint4QuantizeFunc quantize_func_{nullptr}; +}; + +INDEX_FACTORY_REGISTER_REFORMER_ALIAS(UniformUint4Reformer, + UniformUint4Reformer, + 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 827312412..ff334b263 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -103,6 +103,8 @@ enum class QuantizerType { kUniformUint7 = 8, // Global uniform quantization with the full uint8 code range [0, 255]. kUniformUint8 = 9, + // Global uniform quantization with packed 4-bit codes in [0, 15]. + kUniformUint4 = 10, }; struct ZVEC_CORE_API SerializableBase { diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 1c90ce2d4..1967caf89 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -56,6 +56,12 @@ using QueryPreprocessFunc = using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale, float bias, int8_t *out); +// Packed global uint4 quantization. Two codes are stored per byte (low nibble +// first), and the logical dimension is padded to a multiple of 128. +using UniformUint4QuantizeFunc = void (*)(const float *in, size_t dim, + float minimum, float range, + uint8_t *out); + // Direct FP32 conversion. The output layout is selected by get_convert_func(). using ConvertFunc = void (*)(const float *in, size_t dim, void *out); @@ -139,6 +145,7 @@ enum class QuantizeType { // physical representation. Used for kernel dispatch; no serialized // quantizer payload is required. kRaw = 8, + kUniformUint4 = 9, // Uniform uint4: two packed codes per byte. }; enum class RotateType : uint16_t { @@ -197,6 +204,10 @@ ZVEC_TURBO_API DistanceKernels get_distance_kernels( ZVEC_TURBO_API UniformQuantizeFunc get_uniform_quantize_func(DataType data_type); +// Returns the SIMD packed uint4 quantizer, or nullptr when unavailable. +ZVEC_TURBO_API UniformUint4QuantizeFunc +get_uniform_uint4_quantize_func(DataType data_type); + // Returns an optimized fp32 conversion kernel for the requested physical // target type, or nullptr when no optimized implementation is available. // Currently kFp16 and kUint8 are supported. diff --git a/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc b/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc new file mode 100644 index 000000000..7e711f68b --- /dev/null +++ b/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc @@ -0,0 +1,50 @@ +// Copyright 2025-present the zvec project +// SPDX-License-Identifier: Apache-2.0 + +#include "avx512_vnni/uniform_uint4/quantize.h" +#include +#include +#include + +namespace zvec::turbo::avx512_vnni { + +void uniform_uint4_quantize(const float *input, std::size_t dimension, + float minimum, float range, std::uint8_t *output) { + const std::size_t encoded_dimension = ((dimension + 127U) / 128U * 128U) / 2U; + std::memset(output, 0, encoded_dimension); + + constexpr float kAlmostHalf = 0.4999999701976776123046875f; + const __m512 min_value = _mm512_set1_ps(minimum); + const __m512 range_value = _mm512_set1_ps(range); + const __m512 zero = _mm512_setzero_ps(); + const __m512 one = _mm512_set1_ps(1.0f); + const __m512 levels = _mm512_set1_ps(15.0f); + const __m512 almost_half = _mm512_set1_ps(kAlmostHalf); + + alignas(64) int32_t codes[16]; + std::size_t d = 0; + for (; d + 16 <= dimension; d += 16) { + __m512 values = _mm512_loadu_ps(input + d); + values = _mm512_div_ps(_mm512_sub_ps(values, min_value), range_value); + values = _mm512_min_ps(one, _mm512_max_ps(zero, values)); + values = _mm512_add_ps(_mm512_mul_ps(values, levels), almost_half); + _mm512_store_si512(codes, _mm512_cvttps_epi32(values)); + for (std::size_t lane = 0; lane < 16; lane += 2) { + output[(d + lane) >> 1U] = static_cast( + codes[lane] | (static_cast(codes[lane + 1]) << 4U)); + } + } + for (; d < dimension; ++d) { + float normalized = (input[d] - minimum) / range; + normalized = std::min(1.0f, std::max(0.0f, normalized)); + const auto code = static_cast( + static_cast(normalized * 15.0f + kAlmostHalf)); + if ((d & 1U) == 0) { + output[d >> 1U] = code; + } else { + output[d >> 1U] |= static_cast(code << 4U); + } + } +} + +} // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.h b/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.h new file mode 100644 index 000000000..9fc18dbb3 --- /dev/null +++ b/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.h @@ -0,0 +1,13 @@ +// Copyright 2025-present the zvec project +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include + +namespace zvec::turbo::avx512_vnni { + +void uniform_uint4_quantize(const float *input, std::size_t dimension, + float minimum, float range, std::uint8_t *output); + +} // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc new file mode 100644 index 000000000..c820f0b31 --- /dev/null +++ b/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc @@ -0,0 +1,109 @@ +// Copyright 2025-present the zvec project +// SPDX-License-Identifier: Apache-2.0 + +#include "avx512_vnni/uniform_uint4/squared_euclidean.h" +#include +#include +#include "zvec/ailego/internal/platform.h" + +namespace zvec::turbo::avx512_vnni { +namespace { + +inline int32_t Reduce(__m512i value) { + return _mm512_reduce_add_epi32(value); +} + +inline __m512i Accumulate(__m512i sum, __m512i packed, __m512i query_low, + __m512i query_high, __m512i nibble_mask) { + const __m512i low = _mm512_and_si512(packed, nibble_mask); + const __m512i high = + _mm512_and_si512(_mm512_srli_epi16(packed, 4), nibble_mask); + const __m512i low_delta = _mm512_abs_epi8(_mm512_sub_epi8(low, query_low)); + const __m512i high_delta = _mm512_abs_epi8(_mm512_sub_epi8(high, query_high)); + sum = _mm512_dpbusd_epi32(sum, low_delta, low_delta); + return _mm512_dpbusd_epi32(sum, high_delta, high_delta); +} + +static ailego_force_inline void Distance(const uint8_t *lhs, const uint8_t *rhs, + size_t encoded_dimension, + float *distance) { + const __m512i mask = _mm512_set1_epi8(0x0f); + __m512i sum = _mm512_setzero_si512(); + size_t offset = 0; + for (; offset + 64 <= encoded_dimension; offset += 64) { + const __m512i query = _mm512_loadu_si512(rhs + offset); + const __m512i query_low = _mm512_and_si512(query, mask); + const __m512i query_high = + _mm512_and_si512(_mm512_srli_epi16(query, 4), mask); + sum = Accumulate(sum, _mm512_loadu_si512(lhs + offset), query_low, + query_high, mask); + } + int64_t total = Reduce(sum); + for (; offset < encoded_dimension; ++offset) { + const int low_delta = static_cast(lhs[offset] & 0x0fU) - + static_cast(rhs[offset] & 0x0fU); + const int high_delta = static_cast(lhs[offset] >> 4U) - + static_cast(rhs[offset] >> 4U); + total += low_delta * low_delta + high_delta * high_delta; + } + *distance = static_cast(total); +} + +static ailego_force_inline void DistanceFour(const void *const *vectors, + const uint8_t *query, + size_t encoded_dimension, + float *distances) { + const __m512i mask = _mm512_set1_epi8(0x0f); + __m512i sums[4] = {_mm512_setzero_si512(), _mm512_setzero_si512(), + _mm512_setzero_si512(), _mm512_setzero_si512()}; + size_t offset = 0; + for (; offset + 64 <= encoded_dimension; offset += 64) { + const __m512i packed_query = _mm512_loadu_si512(query + offset); + const __m512i query_low = _mm512_and_si512(packed_query, mask); + const __m512i query_high = + _mm512_and_si512(_mm512_srli_epi16(packed_query, 4), mask); + for (size_t lane = 0; lane < 4; ++lane) { + const auto *row = static_cast(vectors[lane]); + sums[lane] = Accumulate(sums[lane], _mm512_loadu_si512(row + offset), + query_low, query_high, mask); + } + } + for (size_t lane = 0; lane < 4; ++lane) { + distances[lane] = static_cast(Reduce(sums[lane])); + } + if (offset < encoded_dimension) { + for (size_t lane = 0; lane < 4; ++lane) { + float tail = 0.0f; + Distance(static_cast(vectors[lane]) + offset, + query + offset, encoded_dimension - offset, &tail); + distances[lane] += tail; + } + } +} + +} // namespace + +void uniform_squared_euclidean_uint4_distance(const void *lhs, const void *rhs, + size_t encoded_dimension, + float *distance) { + Distance(static_cast(lhs), static_cast(rhs), + encoded_dimension, distance); +} + +void uniform_squared_euclidean_uint4_batch_distance(const void *const *vectors, + const void *query, + size_t count, + size_t encoded_dimension, + float *distances) { + const auto *packed_query = static_cast(query); + size_t i = 0; + for (; i + 4 <= count; i += 4) { + DistanceFour(vectors + i, packed_query, encoded_dimension, distances + i); + } + for (; i < count; ++i) { + Distance(static_cast(vectors[i]), packed_query, + encoded_dimension, distances + i); + } +} + +} // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.h b/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.h new file mode 100644 index 000000000..331c0bd24 --- /dev/null +++ b/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.h @@ -0,0 +1,20 @@ +// Copyright 2025-present the zvec project +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include + +namespace zvec::turbo::avx512_vnni { + +// `dimension` is the encoded byte count. Each byte stores two unsigned +// four-bit codes, with the low nibble first. +void uniform_squared_euclidean_uint4_distance(const void *lhs, const void *rhs, + std::size_t dimension, + float *distance); +void uniform_squared_euclidean_uint4_batch_distance(const void *const *vectors, + const void *query, + std::size_t count, + std::size_t dimension, + float *distances); + +} // namespace zvec::turbo::avx512_vnni diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index d70619a2e..afe543614 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -23,6 +23,8 @@ #include "avx512_vnni/raw_uint8/squared_euclidean.h" #include "avx512_vnni/record_quantized_int8/cosine.h" #include "avx512_vnni/record_quantized_int8/squared_euclidean.h" +#include "avx512_vnni/uniform_uint4/quantize.h" +#include "avx512_vnni/uniform_uint4/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" @@ -182,6 +184,12 @@ constexpr KernelSet kKernelTable[] = { avx512_vnni::uniform_squared_euclidean_uint8_batch_distance, avx512_vnni::uniform_squared_euclidean_uint8_query_preprocess}, + // --- uniform-quantized uint4 (packed; AVX512-VNNI only) --- + {QuantizeType::kUniformUint4, DataType::kInt4, CpuArchType::kAVX512VNNI, + MetricType::kSquaredEuclidean, + avx512_vnni::uniform_squared_euclidean_uint4_distance, + avx512_vnni::uniform_squared_euclidean_uint4_batch_distance, nullptr}, + // --- fp16 (scalar) --- {QuantizeType::kFp16, DataType::kFp16, CpuArchType::kScalar, MetricType::kSquaredEuclidean, scalar::squared_euclidean_fp16_distance, @@ -308,6 +316,14 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } +UniformUint4QuantizeFunc get_uniform_uint4_quantize_func(DataType data_type) { + if (data_type == DataType::kInt4 && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { + return avx512_vnni::uniform_uint4_quantize; + } + return nullptr; +} + ConvertFunc get_convert_func(DataType target_data_type) { const ConvertKernel *k = FindConvertKernel(target_data_type); return k ? k->convert : nullptr; diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index afec37637..8fb10f435 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -320,6 +320,8 @@ TEST(IndexInterface, ReopenRestoresUniformReformer) { "test_uniform_uint7_reopen.index"}, {QuantizerType::kUniformUint8, "UniformUint8Converter", "test_uniform_uint8_reopen.index"}, + {QuantizerType::kUniformUint4, "UniformUint4Converter", + "test_uniform_uint4_reopen.index"}, }; for (const auto &test_case : test_cases) { diff --git a/tests/core/metric/uniform_uint4_metric_test.cc b/tests/core/metric/uniform_uint4_metric_test.cc new file mode 100644 index 000000000..05e4a4960 --- /dev/null +++ b/tests/core/metric/uniform_uint4_metric_test.cc @@ -0,0 +1,102 @@ +// Copyright 2025-present the zvec project +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include +#include +#include +#include +#include "metric/metric_params.h" + +namespace zvec::core { +namespace { + +float ScalarDistance(const uint8_t *lhs, const uint8_t *rhs, size_t bytes) { + int64_t sum = 0; + for (size_t i = 0; i < bytes; ++i) { + const int low = + static_cast(lhs[i] & 15U) - static_cast(rhs[i] & 15U); + const int high = + static_cast(lhs[i] >> 4U) - static_cast(rhs[i] >> 4U); + sum += low * low + high * high; + } + return static_cast(sum); +} + +IndexMetric::Pointer CreateMetric(size_t encoded_dimension) { + auto metric = IndexFactory::CreateMetric("UniformUint4"); + if (!metric) return nullptr; + IndexMeta meta(IndexMeta::DataType::DT_INT8, encoded_dimension); + ailego::Params params; + params.set(UNIFORM_UINT4_METRIC_ORIGIN_METRIC_NAME, + std::string("SquaredEuclidean")); + return metric->init(meta, params) == 0 ? metric : nullptr; +} + +TEST(UniformUint4Metric, PairAndBatchMatchScalarExactly) { + std::mt19937 generator(20260807); + std::uniform_int_distribution bytes(0, 255); + for (const size_t logical_dimension : {128UL, 256UL, 1024UL, 65536UL}) { + const size_t encoded_dimension = logical_dimension / 2U; + auto metric = CreateMetric(encoded_dimension); + ASSERT_NE(nullptr, metric); + auto distance = metric->distance(); + auto batch_distance = metric->batch_distance(); + ASSERT_TRUE(static_cast(distance)); + ASSERT_TRUE(static_cast(batch_distance)); + + constexpr size_t count = 7; + std::vector query(encoded_dimension); + std::vector> rows( + count, std::vector(encoded_dimension)); + std::vector pointers(count); + std::vector expected(count); + std::vector actual(count); + for (auto &value : query) value = static_cast(bytes(generator)); + for (size_t i = 0; i < count; ++i) { + for (auto &value : rows[i]) + value = static_cast(bytes(generator)); + pointers[i] = rows[i].data(); + expected[i] = + ScalarDistance(rows[i].data(), query.data(), encoded_dimension); + float pair = 0.0f; + distance(rows[i].data(), query.data(), encoded_dimension, &pair); + EXPECT_EQ(expected[i], pair); + } + batch_distance(pointers.data(), query.data(), count, encoded_dimension, + actual.data()); + EXPECT_EQ(expected, actual) << "logical_dimension=" << logical_dimension; + } +} + +TEST(UniformUint4Metric, QuantizeMatchesReimplPackingAndPadding) { + auto quantize = + turbo::get_uniform_uint4_quantize_func(turbo::DataType::kInt4); + if (!quantize) GTEST_SKIP() << "AVX-512 VNNI is unavailable"; + + constexpr size_t dimension = 131; + constexpr size_t encoded_dimension = 128; + std::vector input(dimension); + for (size_t i = 0; i < dimension; ++i) { + input[i] = -5.0f + static_cast(i % 31U) * 0.5f; + } + std::vector actual(encoded_dimension, 0xff); + std::vector expected(encoded_dimension, 0); + constexpr float minimum = -3.25f; + constexpr float range = 11.5f; + constexpr float almost_half = 0.4999999701976776123046875f; + for (size_t d = 0; d < dimension; ++d) { + float normalized = (input[d] - minimum) / range; + normalized = std::min(1.0f, std::max(0.0f, normalized)); + const auto code = static_cast( + static_cast(normalized * 15.0f + almost_half)); + expected[d >> 1U] |= static_cast(code << (4U * (d & 1U))); + } + quantize(input.data(), dimension, minimum, range, actual.data()); + EXPECT_EQ(expected, actual); +} + +} // namespace +} // namespace zvec::core diff --git a/tests/core/quantizer/uniform_uint4_reformer_test.cc b/tests/core/quantizer/uniform_uint4_reformer_test.cc new file mode 100644 index 000000000..a081907db --- /dev/null +++ b/tests/core/quantizer/uniform_uint4_reformer_test.cc @@ -0,0 +1,154 @@ +// Copyright 2025-present the zvec project +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "quantizer/quantizer_params.h" + +namespace zvec::core { +namespace { + +std::vector ScalarEncode(const float *input, size_t dimension, + float minimum, float range) { + const size_t encoded_dimension = ((dimension + 127U) / 128U * 128U) / 2U; + std::vector output(encoded_dimension, 0); + constexpr float almost_half = 0.4999999701976776123046875f; + for (size_t d = 0; d < dimension; ++d) { + float normalized = (input[d] - minimum) / range; + normalized = std::min(1.0f, std::max(0.0f, normalized)); + const auto code = static_cast( + static_cast(normalized * 15.0f + almost_half)); + output[d >> 1U] |= static_cast(code << (4U * (d & 1U))); + } + return output; +} + +TEST(UniformUint4Reformer, ExactClippedCalibrationPackingAndPersistence) { + constexpr size_t count = 100; + constexpr size_t dimension = 3; + constexpr size_t encoded_dimension = 64; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto converter = IndexFactory::CreateConverter("UniformUint4Converter"); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(meta, ailego::Params())); + + auto holder = + std::make_shared>( + dimension); + float next_value = -150.0f; + for (size_t i = 0; i < count; ++i) { + ailego::NumericalVector vector(dimension); + for (size_t d = 0; d < dimension; ++d) vector[d] = next_value++; + holder->emplace(i, vector); + } + ASSERT_EQ(0, IndexConverter::TrainAndTransform(converter, holder)); + + // N=300: tail=floor(float(N)*0.01)+1=4. Therefore the selected + // order-statistic ranks are 3 and 296, matching reimpl/vamana and KGN. + float minimum = 0.0f; + float range = 0.0f; + uint32_t original_dimension = 0; + const auto ¶ms = converter->meta().reformer_params(); + ASSERT_TRUE(params.get(UNIFORM_UINT4_REFORMER_MINIMUM, &minimum)); + ASSERT_TRUE(params.get(UNIFORM_UINT4_REFORMER_RANGE, &range)); + ASSERT_TRUE(params.get(UNIFORM_UINT4_REFORMER_ORIGINAL_DIMENSION, + &original_dimension)); + EXPECT_FLOAT_EQ(-147.0f, minimum); + EXPECT_FLOAT_EQ(293.0f, range); + EXPECT_EQ(dimension, original_dimension); + EXPECT_EQ(encoded_dimension, converter->meta().dimension()); + EXPECT_EQ("UniformUint4", converter->meta().metric_name()); + + auto encoded_holder = converter->result(); + ASSERT_NE(nullptr, encoded_holder); + EXPECT_EQ(IndexMeta::DataType::DT_INT8, encoded_holder->data_type()); + EXPECT_EQ(encoded_dimension, encoded_holder->dimension()); + auto raw_iter = holder->create_iterator(); + auto encoded_iter = encoded_holder->create_iterator(); + ASSERT_TRUE(raw_iter->is_valid()); + ASSERT_TRUE(encoded_iter->is_valid()); + const auto expected = ScalarEncode( + static_cast(raw_iter->data()), dimension, minimum, range); + EXPECT_EQ( + 0, std::memcmp(expected.data(), encoded_iter->data(), encoded_dimension)); + + auto reformer = IndexFactory::CreateReformer("UniformUint4Reformer"); + ASSERT_NE(nullptr, reformer); + ASSERT_EQ(0, reformer->init(params)); + std::string transformed; + IndexQueryMeta transformed_meta; + ASSERT_EQ(0, reformer->transform( + raw_iter->data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, dimension), + &transformed, &transformed_meta)); + EXPECT_EQ(encoded_dimension, transformed_meta.dimension()); + EXPECT_EQ(std::string(reinterpret_cast(expected.data()), + expected.size()), + transformed); +} + +TEST(UniformUint4Reformer, BuildsDirectlyFromNativeFp16Holder) { + constexpr size_t dimension = 3; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto converter = IndexFactory::CreateConverter("UniformUint4Converter"); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(meta, ailego::Params())); + + auto holder = + std::make_shared>( + dimension); + for (size_t i = 0; i < 2; ++i) { + ailego::NumericalVector vector(dimension); + for (size_t d = 0; d < dimension; ++d) { + vector[d] = static_cast(i * dimension + d); + } + ASSERT_TRUE(holder->emplace(i, vector)); + } + ASSERT_EQ(0, IndexConverter::TrainAndTransform(converter, holder)); + + float minimum = 0.0f; + float range = 0.0f; + ASSERT_TRUE(converter->meta().reformer_params().get( + UNIFORM_UINT4_REFORMER_MINIMUM, &minimum)); + ASSERT_TRUE(converter->meta().reformer_params().get( + UNIFORM_UINT4_REFORMER_RANGE, &range)); + EXPECT_FLOAT_EQ(0.0f, minimum); + EXPECT_FLOAT_EQ(5.0f, range); + + auto encoded_iter = converter->result()->create_iterator(); + ASSERT_TRUE(encoded_iter->is_valid()); + const std::vector first{0.0f, 1.0f, 2.0f}; + const auto expected = ScalarEncode(first.data(), dimension, minimum, range); + EXPECT_EQ( + 0, std::memcmp(expected.data(), encoded_iter->data(), expected.size())); + + auto reformer = IndexFactory::CreateReformer("UniformUint4Reformer"); + ASSERT_NE(nullptr, reformer); + ASSERT_EQ(0, reformer->init(converter->meta().reformer_params())); + auto native_iter = holder->create_iterator(); + std::string converted; + IndexQueryMeta converted_meta; + ASSERT_EQ(0, reformer->convert( + native_iter->data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP16, dimension), + &converted, &converted_meta)); + EXPECT_EQ(std::string(reinterpret_cast(expected.data()), + expected.size()), + converted); +} + +} // namespace +} // namespace zvec::core From 7277f9bf993707261ed06d3f3d7417774ca6d5d9 Mon Sep 17 00:00:00 2001 From: Jianning Wang Date: Tue, 25 Aug 2026 15:11:09 +0800 Subject: [PATCH 2/3] fix(turbo): unify datatype --- src/core/quantizer/uniform_uint4_converter.cc | 3 ++- src/core/quantizer/uniform_uint4_reformer.cc | 2 +- src/core/quantizer/uniform_uint7_converter.cc | 2 +- src/core/quantizer/uniform_uint7_reformer.cc | 2 +- src/include/zvec/turbo/turbo.h | 3 +++ src/turbo/turbo.cc | 6 ++++-- tests/core/metric/uniform_uint4_metric_test.cc | 2 +- 7 files changed, 13 insertions(+), 7 deletions(-) diff --git a/src/core/quantizer/uniform_uint4_converter.cc b/src/core/quantizer/uniform_uint4_converter.cc index 57b8eb417..18c6a514b 100644 --- a/src/core/quantizer/uniform_uint4_converter.cc +++ b/src/core/quantizer/uniform_uint4_converter.cc @@ -364,7 +364,8 @@ class UniformUint4Converter : public IndexConverter { minimum_(minimum), range_(range), quantize_func_( - turbo::get_uniform_uint4_quantize_func(turbo::DataType::kInt4)) {} + turbo::get_uniform_uint4_quantize_func(turbo::DataType::kUint4)) { + } size_t count(void) const override { return front_->count(); diff --git a/src/core/quantizer/uniform_uint4_reformer.cc b/src/core/quantizer/uniform_uint4_reformer.cc index d6fc9f8c0..739e283b9 100644 --- a/src/core/quantizer/uniform_uint4_reformer.cc +++ b/src/core/quantizer/uniform_uint4_reformer.cc @@ -72,7 +72,7 @@ class UniformUint4Reformer : public IndexReformer { const float step = range_ / 15.0f; distance_scale_ = step * step; quantize_func_ = - turbo::get_uniform_uint4_quantize_func(turbo::DataType::kInt4); + turbo::get_uniform_uint4_quantize_func(turbo::DataType::kUint4); initialized_ = true; return 0; } diff --git a/src/core/quantizer/uniform_uint7_converter.cc b/src/core/quantizer/uniform_uint7_converter.cc index b58186707..93a6c9844 100644 --- a/src/core/quantizer/uniform_uint7_converter.cc +++ b/src/core/quantizer/uniform_uint7_converter.cc @@ -306,7 +306,7 @@ class UniformUint7Converter : public IndexConverter { scale_(scale), bias_(bias), quantize_func_( - turbo::get_uniform_quantize_func(turbo::DataType::kInt8)) {} + turbo::get_uniform_quantize_func(turbo::DataType::kUint7)) {} size_t count(void) const override { return front_->count(); diff --git a/src/core/quantizer/uniform_uint7_reformer.cc b/src/core/quantizer/uniform_uint7_reformer.cc index 686c24d21..3142e7516 100644 --- a/src/core/quantizer/uniform_uint7_reformer.cc +++ b/src/core/quantizer/uniform_uint7_reformer.cc @@ -64,7 +64,7 @@ class UniformUint7Reformer : public IndexReformer { bias_ = bias; scale_reciprocal_sq_ = 1.0f / (scale_ * scale_); initialized_ = true; - quantize_func_ = turbo::get_uniform_quantize_func(turbo::DataType::kInt8); + quantize_func_ = turbo::get_uniform_quantize_func(turbo::DataType::kUint7); LOG_INFO("UniformUint7Reformer init: scale=%f, bias=%f, simd=%s", scale_, bias_, quantize_func_ != nullptr ? "avx512" : "scalar"); diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 1967caf89..3ba368324 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -129,6 +129,9 @@ enum class DataType { kFp32, kUint8, kUnknown, + kUint4, + kUint7, + kUint8 }; enum class QuantizeType { diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index afe543614..5a5ae5e26 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -305,7 +305,7 @@ QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, } UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { - if (data_type == DataType::kInt8) { + if (data_type == DataType::kUint7) { // Quantize uses AVX-512F (no VNNI required), but we gate on the same // AVX512_VNNI flag for now since the kernel lives in the avx512_vnni // directory and is compiled with the same march flag. @@ -317,7 +317,9 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { } UniformUint4QuantizeFunc get_uniform_uint4_quantize_func(DataType data_type) { - if (data_type == DataType::kInt4 && + // TODO: unify uniform_uint4_quantize/uniform_uint4_quantize param list and + // merge get_uniform_quantize_func/get_uniform_uint4_quantize_func + if (data_type == DataType::kUint4 && zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI) { return avx512_vnni::uniform_uint4_quantize; } diff --git a/tests/core/metric/uniform_uint4_metric_test.cc b/tests/core/metric/uniform_uint4_metric_test.cc index 05e4a4960..23b9934e9 100644 --- a/tests/core/metric/uniform_uint4_metric_test.cc +++ b/tests/core/metric/uniform_uint4_metric_test.cc @@ -73,7 +73,7 @@ TEST(UniformUint4Metric, PairAndBatchMatchScalarExactly) { TEST(UniformUint4Metric, QuantizeMatchesReimplPackingAndPadding) { auto quantize = - turbo::get_uniform_uint4_quantize_func(turbo::DataType::kInt4); + turbo::get_uniform_uint4_quantize_func(turbo::DataType::kUint4); if (!quantize) GTEST_SKIP() << "AVX-512 VNNI is unavailable"; constexpr size_t dimension = 131; From ed8c3a3769563469092c893212c5e7e239ae5f46 Mon Sep 17 00:00:00 2001 From: Jianning Wang Date: Tue, 25 Aug 2026 15:24:57 +0800 Subject: [PATCH 3/3] fix(turbo): compilation guard --- src/include/zvec/turbo/turbo.h | 1 - .../avx512_vnni/uniform_uint4/quantize.cc | 14 +++++++++++++ .../uniform_uint4/squared_euclidean.cc | 21 ++++++++++++++++++- 3 files changed, 34 insertions(+), 2 deletions(-) diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 3ba368324..cb10a7918 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -131,7 +131,6 @@ enum class DataType { kUnknown, kUint4, kUint7, - kUint8 }; enum class QuantizeType { diff --git a/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc b/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc index 7e711f68b..7ae131ac2 100644 --- a/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc +++ b/src/turbo/distance/avx512_vnni/uniform_uint4/quantize.cc @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 #include "avx512_vnni/uniform_uint4/quantize.h" + +#if defined(__AVX512F__) || (defined(_MSC_VER) && defined(__AVX512F__)) #include #include #include @@ -48,3 +50,15 @@ void uniform_uint4_quantize(const float *input, std::size_t dimension, } } // namespace zvec::turbo::avx512_vnni + +#else // no AVX-512 support + +namespace zvec::turbo::avx512_vnni { + +void uniform_uint4_quantize(const float * /*input*/, std::size_t /*dimension*/, + float /*minimum*/, float /*range*/, + std::uint8_t * /*output*/) {} + +} // namespace zvec::turbo::avx512_vnni + +#endif diff --git a/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc b/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc index c820f0b31..4c126666e 100644 --- a/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc +++ b/src/turbo/distance/avx512_vnni/uniform_uint4/squared_euclidean.cc @@ -2,9 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 #include "avx512_vnni/uniform_uint4/squared_euclidean.h" +#include "zvec/ailego/internal/platform.h" + +#if defined(__AVX512VNNI__) || (defined(_MSC_VER) && defined(__AVX512F__)) #include #include -#include "zvec/ailego/internal/platform.h" namespace zvec::turbo::avx512_vnni { namespace { @@ -107,3 +109,20 @@ void uniform_squared_euclidean_uint4_batch_distance(const void *const *vectors, } } // namespace zvec::turbo::avx512_vnni + +#else // no AVX512-VNNI support + +namespace zvec::turbo::avx512_vnni { + +void uniform_squared_euclidean_uint4_distance(const void * /*lhs*/, + const void * /*rhs*/, + size_t /*encoded_dimension*/, + float * /*distance*/) {} + +void uniform_squared_euclidean_uint4_batch_distance( + const void *const * /*vectors*/, const void * /*query*/, size_t /*count*/, + size_t /*encoded_dimension*/, float * /*distances*/) {} + +} // namespace zvec::turbo::avx512_vnni + +#endif