diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 73dd7950d..0bdfe39e9 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -290,7 +290,11 @@ int Index::Init(const BaseIndexParam ¶m) { return core::IndexError_Runtime; } - if (CreateAndInitConverterReformer(param.quantizer_param, param) != 0) { + // an absent quantizer param behaves the same as a kNone one + const auto quantizer_param = param.quantizer_param + ? param.quantizer_param + : std::make_shared(); + if (CreateAndInitConverterReformer(*quantizer_param, param) != 0) { LOG_ERROR("Failed to create and init converter"); return core::IndexError_Runtime; } diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc index 29ce8f4cc..122e699c5 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -185,9 +185,15 @@ ailego::JsonObject BaseIndexParam::SerializeToJsonObject( // if (preprocess_param) { // json.set("preprocess_param", preprocess_param->SerializeToJson()); // } - if (!omit_empty_value || quantizer_param.type != QuantizerType::kNone) { + if (quantizer_param) { + if (!omit_empty_value || quantizer_param->type != QuantizerType::kNone) { + json_obj.set("quantizer_param", + quantizer_param->SerializeToJsonObject(omit_empty_value)); + } + } else if (!omit_empty_value) { + // no quantizer configured, keep the default(kNone) object as before json_obj.set("quantizer_param", - quantizer_param.SerializeToJsonObject(omit_empty_value)); + QuantizerParam().SerializeToJsonObject(false)); } // if (refiner_param) { // json.set("refiner_param", refiner_param->SerializeToJson()); @@ -239,7 +245,20 @@ bool BaseIndexParam::DeserializeFromJsonObject( if (json_obj.has("quantizer_param")) { if (json_obj.get("quantizer_param", &tmp_json_value); tmp_json_value.is_object()) { - quantizer_param.DeserializeFromJsonObject(tmp_json_value.as_object()); + const auto &quantizer_json_obj = tmp_json_value.as_object(); + // the concrete param type is determined by the quantizer type + auto quantizer_type = QuantizerType::kNone; + ailego::JsonValue type_json_value; + if (!extract_enum_from_json( + quantizer_json_obj, "type", quantizer_type, type_json_value)) { + LOG_ERROR("Error when deserialize json - field:quantizer_param.type"); + return false; + } + auto quantizer = QuantizerParam::Create(quantizer_type); + if (!quantizer->DeserializeFromJsonObject(quantizer_json_obj)) { + LOG_ERROR("Error when deserialize json - field:quantizer_param"); + } + quantizer_param = std::move(quantizer); } } @@ -393,6 +412,33 @@ bool QuantizerParam::DeserializeFromJsonObject( return true; } +QuantizerParam::Pointer QuantizerParam::Create(QuantizerType t) { + switch (t) { + case QuantizerType::kPQ: + return std::make_shared(); + default: + return std::make_shared(t); + } +} + +ailego::JsonObject PqQuantizerParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = QuantizerParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("num_chunk", ailego::JsonValue(num_chunk)); + json_obj.set("num_bits", ailego::JsonValue(num_bits)); + return json_obj; +} + +bool PqQuantizerParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!QuantizerParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + DESERIALIZE_VALUE_FIELD(json_obj, num_chunk); + DESERIALIZE_VALUE_FIELD(json_obj, num_bits); + return true; +} + // bool BaseIndexQueryParam::DeserializeFromJsonObject( // const ailego::JsonObject &json_obj) { // DESERIALIZE_ENUM_FIELD(json_obj, index_type, IndexType); diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index dec1177cc..1685fb443 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -357,7 +357,7 @@ class ProximaEngineHelper { convert_to_engine_quantize_type(db_index_params->quantize_type()); quantize_type.has_value()) { index_param_builder->WithQuantizerParam( - core_interface::QuantizerParam(quantize_type.value())); + core_interface::QuantizerParam::Create(quantize_type.value())); } else { return tl::make_unexpected( Status::InvalidArgument("unsupported quantize type")); diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 877e6e3a0..7a02e1972 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -119,23 +119,49 @@ struct ZVEC_CORE_API SerializableBase { const ailego::JsonObject &json_obj) = 0; }; -// TODO: maybe a base class for quantizer? +//! Common quantizer params shared by all quantizer types struct ZVEC_CORE_API QuantizerParam : public SerializableBase { + using Pointer = std::shared_ptr; + QuantizerType type = QuantizerType::kNone; - int num_subquantizers = 8; // M - int num_bits = 8; // bits per subquantizer bool enable_rotate = false; // rotate vectors before quantization to reduce error // Constructors - // QuantizerParam() = default; - QuantizerParam(QuantizerType t = QuantizerType::kNone, int subquantizers = 8, - int bits = 8, bool rotate = false) - : type(t), - num_subquantizers(subquantizers), - num_bits(bits), - enable_rotate(rotate) {} + QuantizerParam(QuantizerType t = QuantizerType::kNone, bool rotate = false) + : type(t), enable_rotate(rotate) {} + virtual ~QuantizerParam() = default; + + //! Duplicate the param object, keeping the concrete type + virtual Pointer Clone() const { + return std::make_shared(*this); + } + + //! Create the param object matching the quantizer type + static Pointer Create(QuantizerType t); + + protected: + friend class BaseIndexParam; + ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; + + bool DeserializeFromJsonObject(const ailego::JsonObject &json_obj) override; +}; + +//! Product-Quantization specific params +struct PqQuantizerParam : public QuantizerParam { + int num_chunk = 8; // M: number of sub-quantizers + int num_bits = 8; // bits per sub-quantizer + + // Constructors + PqQuantizerParam(int chunks = 8, int bits = 8, bool rotate = false) + : QuantizerParam(QuantizerType::kPQ, rotate), + num_chunk(chunks), + num_bits(bits) {} + QuantizerParam::Pointer Clone() const override { + return std::make_shared(*this); + } protected: friend class BaseIndexParam; @@ -306,7 +332,16 @@ class ZVEC_CORE_API BaseIndexParam : public SerializableBase { // pipeline PreprocessorParam preprocess_param; - QuantizerParam quantizer_param; + //! nullptr means no quantizer is configured (equivalent to kNone) + QuantizerParam::Pointer quantizer_param{nullptr}; + + QuantizerType quantizer_type() const { + return quantizer_param ? quantizer_param->type : QuantizerType::kNone; + } + + bool enable_rotate() const { + return quantizer_param && quantizer_param->enable_rotate; + } BaseIndexQueryParam::Pointer default_query_param = nullptr; // virtual std::shared_ptr GetDefaultQueryParam() const diff --git a/src/include/zvec/core/interface/index_param_builders.h b/src/include/zvec/core/interface/index_param_builders.h index d88057a93..37d77c93a 100644 --- a/src/include/zvec/core/interface/index_param_builders.h +++ b/src/include/zvec/core/interface/index_param_builders.h @@ -59,7 +59,13 @@ class BaseIndexParamBuilder { // : public } ActualIndexParamBuilderType &WithQuantizerParam( const QuantizerParam &quantizer_param) { - param->quantizer_param = quantizer_param; + param->quantizer_param = quantizer_param.Clone(); + return static_cast(*this); + } + ActualIndexParamBuilderType &WithQuantizerParam( + const QuantizerParam::Pointer &quantizer_param) { + param->quantizer_param = + quantizer_param ? quantizer_param->Clone() : nullptr; return static_cast(*this); } // ActualIndexParamBuilderType &WithRefinerParam( @@ -88,7 +94,12 @@ class BaseIndexParamBuilder { // : public } ActualIndexParamBuilderType &WithEnableRotate(bool enable_rotate) { - param->quantizer_param.enable_rotate = enable_rotate; + // copy-on-write: never mutate a param object shared with others + auto quantizer_param = param->quantizer_param + ? param->quantizer_param->Clone() + : std::make_shared(); + quantizer_param->enable_rotate = enable_rotate; + param->quantizer_param = std::move(quantizer_param); return static_cast(*this); } diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 5baa8c418..8b33c4902 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -65,12 +65,42 @@ using RotateFunc = void (*)(const float *in, float *out, size_t in_dim, using UnrotateFunc = void (*)(const float *in, float *out, size_t in_dim, size_t out_dim, void *ctx); +// PQ kernel function pointer types. +// +// ADC: LUT look-up distance between a PQ code and a query (via LUT). +// pq_code: [num_chunk] uint8_t +// lut: [num_chunk * 256] float +// Uses void* to match DistanceFunc signature for direct assignment. +using PqAdcDistanceFunc = void (*)(const void *pq_code, const void *lut, + size_t num_chunk, float *out); + +// SDC kernel: centroid-to-centroid distance between two PQ codes. +// a, b: [num_chunk] uint8_t +// dist_table: [num_chunk * 256 * 256] float +// Uses void* for consistency with DistanceFunc / PqAdcDistanceFunc. +using PqSdcKernelFunc = void (*)(const void *a, const void *b, + const void *dist_table, size_t num_chunk, + float *out); + +// Batch ADC: compute distances for multiple PQ codes against a shared LUT. +// Signature matches BatchDistanceFunc for direct assignment (no lambda). +using PqBatchAdcFunc = void (*)(const void **candidates, const void *lut, + size_t num, size_t num_chunk, float *out); + // ISA-dispatched rotate/unrotate kernels. struct RotatorKernels { RotateFunc rotate = nullptr; UnrotateFunc unrotate = nullptr; }; +// data_type selects the code packing layout: +// kInt8: one uint8 per sub-quantizer (256 centroids, stride=256) +struct PqKernels { + PqAdcDistanceFunc adc_distance = nullptr; + PqSdcKernelFunc sdc_distance = nullptr; + PqBatchAdcFunc batch_adc_distance = nullptr; +}; + enum class MetricType { kSquaredEuclidean, kCosine, @@ -142,4 +172,10 @@ get_uniform_quantize_func(DataType data_type); RotatorKernels get_rotator_kernels( RotateType rotate_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); +// Returns all PQ kernels dispatched for the given data_type, quantize_type +// and CPU arch. +PqKernels get_pq_kernels(DataType data_type, + QuantizeType quantize_type = QuantizeType::kPQ, + CpuArchType cpu_arch_type = CpuArchType::kAuto); + } // namespace zvec::turbo diff --git a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..b63d8357d --- /dev/null +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,255 @@ +// 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. + +// This file is compiled with per-file -march=core-avx2 (set in CMakeLists.txt) +// so that AVX2 intrinsics are available. When the build toolchain cannot emit +// AVX2 code, each function falls back to a no-op stub guarded by +// #if defined(__AVX2__). + +#include "avx2/pq_quantizer_int8/pq_distance.h" +#if defined(__AVX2__) +#include +#endif +#include +#include + +namespace zvec::turbo::avx2 { + +#if defined(__AVX2__) +namespace { + +// Horizontal sum of 8 floats in a __m256 register. +inline float horizontal_sum_avx2(__m256 v) { + // High 128 bits + low 128 bits + __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 lo = _mm256_castps256_ps128(v); + __m128 sum128 = _mm_add_ps(lo, hi); + // Shuffle and add: [a+b, c+d, a+b, c+d] + __m128 shuf = _mm_movehdup_ps(sum128); // [b, b, d, d] + __m128 sum64 = _mm_add_ps(sum128, shuf); + // Final: [a+b+c+d, ..., ...] + __m128 shuf32 = _mm_movehl_ps(sum64, sum64); + __m128 sum32 = _mm_add_ss(sum64, shuf32); + return _mm_cvtss_f32(sum32); +} + +} // namespace +#endif + +void pq_adc_int8_distance_avx2(const void *pq_code_v, const void *lut_v, + size_t num_chunk, float *out) { +#if defined(__AVX2__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 8; // AVX2 processes 8 floats at once + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + + __m256 acc = _mm256_setzero_ps(); + + // Base offsets: [0, 256, 512, 768, 1024, 1280, 1536, 1792] + // These represent m * 256 for m = 0..7 within each chunk. + const __m256i base_offsets = _mm256_setr_epi32( + 0, kNumCentroids, 2 * kNumCentroids, 3 * kNumCentroids, 4 * kNumCentroids, + 5 * kNumCentroids, 6 * kNumCentroids, 7 * kNumCentroids); + + size_t m = 0; + + // Main loop: process 8 subquantizers per iteration + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + // Load 8 uint8 codes and zero-extend to int32 + // pq_code[m..m+7] -> 8 int32 indices + __m128i codes_8x8 = + _mm_loadl_epi64(reinterpret_cast(pq_code + m)); + __m256i codes_8x32 = _mm256_cvtepu8_epi32(codes_8x8); + + // Add base offsets: indices[m] = m * 256 + code[m] + __m256i indices = _mm256_add_epi32(codes_8x32, base_offsets); + + // Gather 8 floats from lut using computed indices + // lut_ptr + indices[i] * scale(4 bytes per float) + __m256 gathered = _mm256_i32gather_ps(lut + m * kNumCentroids, indices, 4); + + acc = _mm256_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx2(acc); + + // Scalar leftover: process remaining subquantizers + for (; m < num_chunk; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + + *out = sum; +#else + (void)pq_code_v; + (void)lut_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_sdc_int8_distance_avx2(const void *a_v, const void *b_v, + const void *dist_table_v, size_t num_chunk, + float *out) { +#if defined(__AVX2__) + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 8; + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + __m256 acc = _mm256_setzero_ps(); + + // Base offsets for SDC: m * 65536 (float indices into dist_table) + const __m256i base_offsets = _mm256_setr_epi32( + 0, kTablePerSub, 2 * kTablePerSub, 3 * kTablePerSub, 4 * kTablePerSub, + 5 * kTablePerSub, 6 * kTablePerSub, 7 * kTablePerSub); + + // Multiplier for a[m] * 256 + const __m256i a_multiplier = _mm256_set1_epi32(kNumCentroids); + + size_t m = 0; + + // Main loop: process 8 subquantizers per iteration + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + // Load a[m..m+7] and b[m..m+7], zero-extend to int32 + __m128i a_8x8 = _mm_loadl_epi64(reinterpret_cast(a + m)); + __m128i b_8x8 = _mm_loadl_epi64(reinterpret_cast(b + m)); + __m256i a_8x32 = _mm256_cvtepu8_epi32(a_8x8); + __m256i b_8x32 = _mm256_cvtepu8_epi32(b_8x8); + + // Compute in-lane index: a[m] * 256 + b[m] + k * 65536 (k = lane, 0..7). + // The m * 65536 offset is applied via the gather base pointer below. + __m256i a_shifted = _mm256_mullo_epi32(a_8x32, a_multiplier); + __m256i indices = _mm256_add_epi32(a_shifted, b_8x32); + indices = _mm256_add_epi32(indices, base_offsets); + + // Gather 8 floats from dist_table. The gather base must include the + // per-iteration m * kTablePerSub offset; base_offsets only carries the + // in-lane k * kTablePerSub component (k = 0..7), so gathering from a + // fixed dist_table base would read the wrong subquantizer tables once + // num_chunk > 8 (m >= 8). + __m256 gathered = + _mm256_i32gather_ps(dist_table + m * kTablePerSub, indices, 4); + + acc = _mm256_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx2(acc); + + // Scalar leftover + for (; m < num_chunk; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + + *out = sum; +#else + (void)a_v; + (void)b_v; + (void)dist_table_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_adc_int8_batch_distance_avx2(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_chunk, float *out) { +#if defined(__AVX2__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 8; + constexpr int kBatch = 4; + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + // Base offsets: [0, 256, 512, ..., 7*256] — reused for all candidates. + const __m256i base_offsets = _mm256_setr_epi32( + 0, kNumCentroids, 2 * kNumCentroids, 3 * kNumCentroids, 4 * kNumCentroids, + 5 * kNumCentroids, 6 * kNumCentroids, 7 * kNumCentroids); + + size_t i = 0; + for (; i + kBatch <= num; i += kBatch) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(); + __m256 acc3 = _mm256_setzero_ps(); + + size_t m = 0; + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + const float *lut_base = lut + m * kNumCentroids; + + __m128i codes0 = + _mm_loadl_epi64(reinterpret_cast(c0 + m)); + __m128i codes1 = + _mm_loadl_epi64(reinterpret_cast(c1 + m)); + __m128i codes2 = + _mm_loadl_epi64(reinterpret_cast(c2 + m)); + __m128i codes3 = + _mm_loadl_epi64(reinterpret_cast(c3 + m)); + + __m256i idx0 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes0), base_offsets); + __m256i idx1 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes1), base_offsets); + __m256i idx2 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes2), base_offsets); + __m256i idx3 = + _mm256_add_epi32(_mm256_cvtepu8_epi32(codes3), base_offsets); + + acc0 = _mm256_add_ps(acc0, _mm256_i32gather_ps(lut_base, idx0, 4)); + acc1 = _mm256_add_ps(acc1, _mm256_i32gather_ps(lut_base, idx1, 4)); + acc2 = _mm256_add_ps(acc2, _mm256_i32gather_ps(lut_base, idx2, 4)); + acc3 = _mm256_add_ps(acc3, _mm256_i32gather_ps(lut_base, idx3, 4)); + } + + float s0 = horizontal_sum_avx2(acc0); + float s1 = horizontal_sum_avx2(acc1); + float s2 = horizontal_sum_avx2(acc2); + float s3 = horizontal_sum_avx2(acc3); + + // Scalar leftover for remaining subquantizers. + for (; m < num_chunk; ++m) { + const float *tab = lut + m * kNumCentroids; + s0 += tab[c0[m]]; + s1 += tab[c1[m]]; + s2 += tab[c2[m]]; + s3 += tab[c3[m]]; + } + out[i] = s0; + out[i + 1] = s1; + out[i + 2] = s2; + out[i + 3] = s3; + } + // Remaining candidates: use single ADC kernel. + for (; i < num; ++i) { + pq_adc_int8_distance_avx2(candidates[i], lut, num_chunk, out + i); + } +#else + (void)candidates_v; + (void)lut_v; + (void)num; + (void)num_chunk; + (void)out; +#endif +} + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..f95aba404 --- /dev/null +++ b/src/turbo/distance/avx2/pq_quantizer_int8/pq_distance.h @@ -0,0 +1,41 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::avx2 { + +// ADC (Asymmetric Distance Computation) via AVX2 gather. +// Processes 8 subquantizers per _mm256_i32gather_ps iteration. +// For general M: loop in chunks of 8, scalar leftover. +void pq_adc_int8_distance_avx2(const void *pq_code, const void *lut, + size_t num_chunk, float *out); + +// SDC (Symmetric Distance Computation) via AVX2 gather. +// Computes indices (a[m]*256 + b[m]) as int32, adds per-subquantizer +// base offsets, gathers 8 floats per iteration. +void pq_sdc_int8_distance_avx2(const void *a, const void *b, + const void *dist_table, size_t num_chunk, + float *out); + +// Batch ADC via AVX2 gather: process 4 candidates per iteration, +// each using 8-wide _mm256_i32gather_ps. 4 independent __m256 +// accumulators maximize ILP. +void pq_adc_int8_batch_distance_avx2(const void **candidates, const void *lut, + size_t num, size_t num_chunk, float *out); + +} // namespace zvec::turbo::avx2 diff --git a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..28fbd1b30 --- /dev/null +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,255 @@ +// 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. + +// This file is compiled with per-file -march=icelake-server (set in +// CMakeLists.txt) so that AVX512 intrinsics are available. When the build +// toolchain cannot emit AVX-512 code, each function falls back to a no-op +// stub guarded by #if defined(__AVX512F__). + +#include "avx512/pq_quantizer_int8/pq_distance.h" +#if defined(__AVX512F__) +#include +#endif +#include +#include + +namespace zvec::turbo::avx512 { + +#if defined(__AVX512F__) +namespace { + +// Horizontal sum of 16 floats in a __m512 register. +inline float horizontal_sum_avx512(__m512 v) { + // Use _mm512_reduce_add_ps which is available in AVX512F + return _mm512_reduce_add_ps(v); +} + +} // namespace +#endif + +void pq_adc_int8_distance_avx512(const void *pq_code_v, const void *lut_v, + size_t num_chunk, float *out) { +#if defined(__AVX512F__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 16; // AVX512 processes 16 floats at once + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + + __m512 acc = _mm512_setzero_ps(); + + // Base offsets: [0, 256, 512, ..., 15*256] = m * 256 for m = 0..15 + const __m512i base_offsets = _mm512_setr_epi32( + 0 * kNumCentroids, 1 * kNumCentroids, 2 * kNumCentroids, + 3 * kNumCentroids, 4 * kNumCentroids, 5 * kNumCentroids, + 6 * kNumCentroids, 7 * kNumCentroids, 8 * kNumCentroids, + 9 * kNumCentroids, 10 * kNumCentroids, 11 * kNumCentroids, + 12 * kNumCentroids, 13 * kNumCentroids, 14 * kNumCentroids, + 15 * kNumCentroids); + + size_t m = 0; + + // Main loop: process 16 subquantizers per iteration + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + // Load 16 uint8 codes and zero-extend to int32 + // Use unaligned load of 16 bytes + __m128i codes_16x8 = + _mm_loadu_si128(reinterpret_cast(pq_code + m)); + __m512i codes_16x32 = _mm512_cvtepu8_epi32(codes_16x8); + + // Add base offsets: indices[m] = m * 256 + code[m] + __m512i indices = _mm512_add_epi32(codes_16x32, base_offsets); + + // Gather 16 floats from lut using computed indices + __m512 gathered = _mm512_i32gather_ps(indices, lut + m * kNumCentroids, 4); + + acc = _mm512_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx512(acc); + + // Scalar leftover: process remaining subquantizers + for (; m < num_chunk; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + + *out = sum; +#else + (void)pq_code_v; + (void)lut_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_sdc_int8_distance_avx512(const void *a_v, const void *b_v, + const void *dist_table_v, size_t num_chunk, + float *out) { +#if defined(__AVX512F__) + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 16; + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + __m512 acc = _mm512_setzero_ps(); + + // Base offsets for SDC: m * 65536 + const __m512i base_offsets = _mm512_setr_epi32( + 0 * kTablePerSub, 1 * kTablePerSub, 2 * kTablePerSub, 3 * kTablePerSub, + 4 * kTablePerSub, 5 * kTablePerSub, 6 * kTablePerSub, 7 * kTablePerSub, + 8 * kTablePerSub, 9 * kTablePerSub, 10 * kTablePerSub, 11 * kTablePerSub, + 12 * kTablePerSub, 13 * kTablePerSub, 14 * kTablePerSub, + 15 * kTablePerSub); + + // Multiplier for a[m] * 256 + const __m512i a_multiplier = _mm512_set1_epi32(kNumCentroids); + + size_t m = 0; + + // Main loop: process 16 subquantizers per iteration + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + // Load a[m..m+15] and b[m..m+15], zero-extend to int32 + __m128i a_16x8 = _mm_loadu_si128(reinterpret_cast(a + m)); + __m128i b_16x8 = _mm_loadu_si128(reinterpret_cast(b + m)); + __m512i a_16x32 = _mm512_cvtepu8_epi32(a_16x8); + __m512i b_16x32 = _mm512_cvtepu8_epi32(b_16x8); + + // Compute in-lane index: a[m] * 256 + b[m] + k * 65536 (k = lane, 0..15). + // The m * 65536 offset is applied via the gather base pointer below. + __m512i a_shifted = _mm512_mullo_epi32(a_16x32, a_multiplier); + __m512i indices = _mm512_add_epi32(a_shifted, b_16x32); + indices = _mm512_add_epi32(indices, base_offsets); + + // Gather 16 floats from dist_table. The gather base must include the + // per-iteration m * kTablePerSub offset; base_offsets only carries the + // in-lane k * kTablePerSub component (k = 0..15), so gathering from a + // fixed dist_table base would read the wrong subquantizer tables once + // num_chunk > 16 (m >= 16). + __m512 gathered = + _mm512_i32gather_ps(indices, dist_table + m * kTablePerSub, 4); + + acc = _mm512_add_ps(acc, gathered); + } + + float sum = horizontal_sum_avx512(acc); + + // Scalar leftover + for (; m < num_chunk; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + + *out = sum; +#else + (void)a_v; + (void)b_v; + (void)dist_table_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_adc_int8_batch_distance_avx512(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_chunk, float *out) { +#if defined(__AVX512F__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 16; + constexpr int kBatch = 4; + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + // Base offsets: [0, 256, 512, ..., 15*256] — reused for all candidates. + const __m512i base_offsets = _mm512_setr_epi32( + 0 * kNumCentroids, 1 * kNumCentroids, 2 * kNumCentroids, + 3 * kNumCentroids, 4 * kNumCentroids, 5 * kNumCentroids, + 6 * kNumCentroids, 7 * kNumCentroids, 8 * kNumCentroids, + 9 * kNumCentroids, 10 * kNumCentroids, 11 * kNumCentroids, + 12 * kNumCentroids, 13 * kNumCentroids, 14 * kNumCentroids, + 15 * kNumCentroids); + + size_t i = 0; + for (; i + kBatch <= num; i += kBatch) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + __m512 acc0 = _mm512_setzero_ps(); + __m512 acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(); + __m512 acc3 = _mm512_setzero_ps(); + + size_t m = 0; + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + const float *lut_base = lut + m * kNumCentroids; + + __m128i codes0 = + _mm_loadu_si128(reinterpret_cast(c0 + m)); + __m128i codes1 = + _mm_loadu_si128(reinterpret_cast(c1 + m)); + __m128i codes2 = + _mm_loadu_si128(reinterpret_cast(c2 + m)); + __m128i codes3 = + _mm_loadu_si128(reinterpret_cast(c3 + m)); + + __m512i idx0 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes0), base_offsets); + __m512i idx1 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes1), base_offsets); + __m512i idx2 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes2), base_offsets); + __m512i idx3 = + _mm512_add_epi32(_mm512_cvtepu8_epi32(codes3), base_offsets); + + acc0 = _mm512_add_ps(acc0, _mm512_i32gather_ps(idx0, lut_base, 4)); + acc1 = _mm512_add_ps(acc1, _mm512_i32gather_ps(idx1, lut_base, 4)); + acc2 = _mm512_add_ps(acc2, _mm512_i32gather_ps(idx2, lut_base, 4)); + acc3 = _mm512_add_ps(acc3, _mm512_i32gather_ps(idx3, lut_base, 4)); + } + + float s0 = _mm512_reduce_add_ps(acc0); + float s1 = _mm512_reduce_add_ps(acc1); + float s2 = _mm512_reduce_add_ps(acc2); + float s3 = _mm512_reduce_add_ps(acc3); + + // Scalar leftover for remaining subquantizers. + for (; m < num_chunk; ++m) { + const float *tab = lut + m * kNumCentroids; + s0 += tab[c0[m]]; + s1 += tab[c1[m]]; + s2 += tab[c2[m]]; + s3 += tab[c3[m]]; + } + out[i] = s0; + out[i + 1] = s1; + out[i + 2] = s2; + out[i + 3] = s3; + } + // Remaining candidates: use single ADC kernel. + for (; i < num; ++i) { + pq_adc_int8_distance_avx512(candidates[i], lut, num_chunk, out + i); + } +#else + (void)candidates_v; + (void)lut_v; + (void)num; + (void)num_chunk; + (void)out; +#endif +} + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..64b35c017 --- /dev/null +++ b/src/turbo/distance/avx512/pq_quantizer_int8/pq_distance.h @@ -0,0 +1,40 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::avx512 { + +// ADC (Asymmetric Distance Computation) via AVX512 gather. +// Processes 16 subquantizers per _mm512_i32gather_ps iteration. +void pq_adc_int8_distance_avx512(const void *pq_code, const void *lut, + size_t num_chunk, float *out); + +// SDC (Symmetric Distance Computation) via AVX512 gather. +// 16-wide index computation + gather. +void pq_sdc_int8_distance_avx512(const void *a, const void *b, + const void *dist_table, size_t num_chunk, + float *out); + +// Batch ADC via AVX512 gather: process 4 candidates per iteration, +// each using 16-wide _mm512_i32gather_ps. 4 independent __m512 +// accumulators maximize ILP. +void pq_adc_int8_batch_distance_avx512(const void **candidates, const void *lut, + size_t num, size_t num_chunk, + float *out); + +} // namespace zvec::turbo::avx512 diff --git a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..9cc313f8f --- /dev/null +++ b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,207 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// NEON lacks hardware gather instructions (like x86 _mm256_i32gather_ps), +// so LUT lookups remain scalar. NEON accelerates the accumulation path: +// 4 floats are loaded into a float32x4_t register and pairwise-added into +// the running accumulator, halving the number of FP add operations. + +#include "neon/pq_quantizer_int8/pq_distance.h" +#if defined(__ARM_NEON) && defined(__aarch64__) +#include +#endif +#include +#include + +namespace zvec::turbo::neon { + +#if defined(__ARM_NEON) && defined(__aarch64__) +namespace { + +// Horizontal sum of 4 floats in a float32x4_t register via pairwise add. +// [a, b, c, d] → [a+b, c+d, a+b, c+d] → [a+b+c+d, ...] +inline float horizontal_sum_neon(float32x4_t v) { + float32x2_t lo = vget_low_f32(v); + float32x2_t hi = vget_high_f32(v); + float32x2_t sum2 = vadd_f32(lo, hi); // [a+b, c+d] + return vget_lane_f32(vpadd_f32(sum2, sum2), 0); // a+b+c+d +} + +} // namespace +#endif + +void pq_adc_int8_distance_neon(const void *pq_code_v, const void *lut_v, + size_t num_chunk, float *out) { +#if defined(__ARM_NEON) && defined(__aarch64__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 4; // NEON processes 4 floats at once + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + + float32x4_t acc = vdupq_n_f32(0.0f); + + size_t m = 0; + + // Main loop: process 4 subquantizers per iteration. + // Scalar LUT lookups (NEON has no gather), then NEON pairwise accumulation. + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + float d0 = lut[(m + 0) * kNumCentroids + pq_code[m + 0]]; + float d1 = lut[(m + 1) * kNumCentroids + pq_code[m + 1]]; + float d2 = lut[(m + 2) * kNumCentroids + pq_code[m + 2]]; + float d3 = lut[(m + 3) * kNumCentroids + pq_code[m + 3]]; + float32x4_t d = {d0, d1, d2, d3}; + acc = vaddq_f32(acc, d); + } + + float sum = horizontal_sum_neon(acc); + + // Scalar leftover: process remaining subquantizers + for (; m < num_chunk; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + + *out = sum; +#else + (void)pq_code_v; + (void)lut_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_sdc_int8_distance_neon(const void *a_v, const void *b_v, + const void *dist_table_v, size_t num_chunk, + float *out) { +#if defined(__ARM_NEON) && defined(__aarch64__) + constexpr int kNumCentroids = 256; + constexpr int kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + constexpr int kChunkSize = 4; + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + + float32x4_t acc = vdupq_n_f32(0.0f); + + size_t m = 0; + + // Main loop: process 4 subquantizers per iteration. + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + float d0 = dist_table[(m + 0) * kTablePerSub + + static_cast(a[m + 0]) * kNumCentroids + + static_cast(b[m + 0])]; + float d1 = dist_table[(m + 1) * kTablePerSub + + static_cast(a[m + 1]) * kNumCentroids + + static_cast(b[m + 1])]; + float d2 = dist_table[(m + 2) * kTablePerSub + + static_cast(a[m + 2]) * kNumCentroids + + static_cast(b[m + 2])]; + float d3 = dist_table[(m + 3) * kTablePerSub + + static_cast(a[m + 3]) * kNumCentroids + + static_cast(b[m + 3])]; + float32x4_t d = {d0, d1, d2, d3}; + acc = vaddq_f32(acc, d); + } + + float sum = horizontal_sum_neon(acc); + + // Scalar leftover + for (; m < num_chunk; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + + *out = sum; +#else + (void)a_v; + (void)b_v; + (void)dist_table_v; + (void)num_chunk; + (void)out; +#endif +} + +void pq_adc_int8_batch_distance_neon(const void **candidates_v, + const void *lut_v, size_t num, + size_t num_chunk, float *out) { +#if defined(__ARM_NEON) && defined(__aarch64__) + constexpr int kNumCentroids = 256; + constexpr int kChunkSize = 4; + constexpr int kBatch = 4; + const auto *lut = reinterpret_cast(lut_v); + const auto *candidates = + reinterpret_cast(candidates_v); + + size_t i = 0; + for (; i + kBatch <= num; i += kBatch) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + float32x4_t acc0 = vdupq_n_f32(0.0f); + float32x4_t acc1 = vdupq_n_f32(0.0f); + float32x4_t acc2 = vdupq_n_f32(0.0f); + float32x4_t acc3 = vdupq_n_f32(0.0f); + + size_t m = 0; + for (; m + kChunkSize <= num_chunk; m += kChunkSize) { + const float *tab = lut + m * kNumCentroids; + + float32x4_t d0 = {tab[c0[m + 0]], tab[c0[m + 1]], tab[c0[m + 2]], + tab[c0[m + 3]]}; + float32x4_t d1 = {tab[c1[m + 0]], tab[c1[m + 1]], tab[c1[m + 2]], + tab[c1[m + 3]]}; + float32x4_t d2 = {tab[c2[m + 0]], tab[c2[m + 1]], tab[c2[m + 2]], + tab[c2[m + 3]]}; + float32x4_t d3 = {tab[c3[m + 0]], tab[c3[m + 1]], tab[c3[m + 2]], + tab[c3[m + 3]]}; + + acc0 = vaddq_f32(acc0, d0); + acc1 = vaddq_f32(acc1, d1); + acc2 = vaddq_f32(acc2, d2); + acc3 = vaddq_f32(acc3, d3); + } + + float s0 = horizontal_sum_neon(acc0); + float s1 = horizontal_sum_neon(acc1); + float s2 = horizontal_sum_neon(acc2); + float s3 = horizontal_sum_neon(acc3); + + // Scalar leftover for remaining subquantizers. + for (; m < num_chunk; ++m) { + const float *tab = lut + m * kNumCentroids; + s0 += tab[c0[m]]; + s1 += tab[c1[m]]; + s2 += tab[c2[m]]; + s3 += tab[c3[m]]; + } + out[i] = s0; + out[i + 1] = s1; + out[i + 2] = s2; + out[i + 3] = s3; + } + // Remaining candidates: use single ADC kernel. + for (; i < num; ++i) { + pq_adc_int8_distance_neon(candidates[i], lut, num_chunk, out + i); + } +#else + (void)candidates_v; + (void)lut_v; + (void)num; + (void)num_chunk; + (void)out; +#endif +} + +} // namespace zvec::turbo::neon diff --git a/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..d0c0e55d4 --- /dev/null +++ b/src/turbo/distance/neon/pq_quantizer_int8/pq_distance.h @@ -0,0 +1,40 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::neon { + +// ADC (Asymmetric Distance Computation) using NEON vector accumulation. +// LUT lookups are scalar (NEON lacks hardware gather); accumulation uses +// float32x4_t with pairwise horizontal reduction. +void pq_adc_int8_distance_neon(const void *pq_code, const void *lut, + size_t num_chunk, float *out); + +// SDC (Symmetric Distance Computation) via scalar lookup. +// The dist_table is too large (65536 floats per subquantizer) for NEON +// table-lookup instructions, so the implementation is scalar. +void pq_sdc_int8_distance_neon(const void *a, const void *b, + const void *dist_table, size_t num_chunk, + float *out); + +// Batch ADC: compute distances for multiple PQ codes against a shared LUT. +// Processes 4 candidates per iteration with NEON vector accumulation. +void pq_adc_int8_batch_distance_neon(const void **candidates, const void *lut, + size_t num, size_t num_chunk, float *out); + +} // namespace zvec::turbo::neon diff --git a/src/turbo/distance/scalar/fp16/cosine.cc b/src/turbo/distance/scalar/fp16/cosine.cc new file mode 100644 index 000000000..a6bd8d61c --- /dev/null +++ b/src/turbo/distance/scalar/fp16/cosine.cc @@ -0,0 +1,40 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp16/cosine.h" +#include "scalar/fp16/inner_product.h" + +namespace zvec::turbo::scalar { + +void cosine_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + // inner_product_fp16_distance returns -real_IP; cosine = 1 - real_IP = 1 + + // ip. + float ip; + inner_product_fp16_distance(a, b, dim, &ip); + + *distance = 1 + ip; +} + +void cosine_fp16_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances) { + inner_product_fp16_batch_distance(vectors, query, n, dim, distances); + // inner_product batch returns -real_IP per element; cosine = 1 - real_IP = + // 1 + d. + for (size_t i = 0; i < n; i++) { + distances[i] = 1 + distances[i]; + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/cosine.h b/src/turbo/distance/scalar/fp16/cosine.h new file mode 100644 index 000000000..39be959b2 --- /dev/null +++ b/src/turbo/distance/scalar/fp16/cosine.h @@ -0,0 +1,30 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute cosine distance between a single FP16 vector pair. +// Returns 1 - dot(a, b). +void cosine_fp16_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of cosine_fp16_distance. +void cosine_fp16_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/inner_product.cc b/src/turbo/distance/scalar/fp16/inner_product.cc new file mode 100644 index 000000000..a8d944b0e --- /dev/null +++ b/src/turbo/distance/scalar/fp16/inner_product.cc @@ -0,0 +1,41 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp16/inner_product.h" +#include + +namespace zvec::turbo::scalar { + +void inner_product_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + const ailego::Float16 *m = reinterpret_cast(a); + const ailego::Float16 *q = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(m[i]) * static_cast(q[i]); + } + + *distance = -sum; +} + +void inner_product_fp16_batch_distance(const void *const *vectors, + const void *query, size_t n, size_t dim, + float *distances) { + for (size_t i = 0; i < n; ++i) { + inner_product_fp16_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/inner_product.h b/src/turbo/distance/scalar/fp16/inner_product.h new file mode 100644 index 000000000..08059ed3d --- /dev/null +++ b/src/turbo/distance/scalar/fp16/inner_product.h @@ -0,0 +1,30 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute negated inner product between a single FP16 vector pair. +void inner_product_fp16_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of inner_product_fp16_distance. +void inner_product_fp16_batch_distance(const void *const *vectors, + const void *query, size_t n, size_t dim, + float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/squared_euclidean.cc b/src/turbo/distance/scalar/fp16/squared_euclidean.cc new file mode 100644 index 000000000..931028b40 --- /dev/null +++ b/src/turbo/distance/scalar/fp16/squared_euclidean.cc @@ -0,0 +1,42 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/fp16/squared_euclidean.h" +#include + +namespace zvec::turbo::scalar { + +void squared_euclidean_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + const ailego::Float16 *m = reinterpret_cast(a); + const ailego::Float16 *q = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + float diff = static_cast(m[i]) - static_cast(q[i]); + sum += diff * diff; + } + + *distance = sum; +} + +void squared_euclidean_fp16_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, float *distances) { + for (size_t i = 0; i < n; ++i) { + squared_euclidean_fp16_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/fp16/squared_euclidean.h b/src/turbo/distance/scalar/fp16/squared_euclidean.h new file mode 100644 index 000000000..b395066e2 --- /dev/null +++ b/src/turbo/distance/scalar/fp16/squared_euclidean.h @@ -0,0 +1,30 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute squared euclidean distance between a single FP16 vector pair. +void squared_euclidean_fp16_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of squared euclidean FP16. +void squared_euclidean_fp16_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc new file mode 100644 index 000000000..6d8ca0d4f --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.cc @@ -0,0 +1,89 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "scalar/pq_quantizer_int8/pq_distance.h" + +namespace zvec::turbo::scalar { + +void pq_adc_int8_distance(const void *pq_code_v, const void *lut_v, + size_t num_chunk, float *out) { + constexpr size_t kNumCentroids = 256; + const auto *pq_code = reinterpret_cast(pq_code_v); + const auto *lut = reinterpret_cast(lut_v); + float sum = 0.0f; + for (size_t m = 0; m < num_chunk; ++m) { + sum += lut[m * kNumCentroids + pq_code[m]]; + } + *out = sum; +} + +void pq_sdc_int8_distance(const void *a_v, const void *b_v, + const void *dist_table_v, size_t num_chunk, + float *out) { + constexpr size_t kNumCentroids = 256; + constexpr size_t kTablePerSub = kNumCentroids * kNumCentroids; // 65536 + const auto *a = reinterpret_cast(a_v); + const auto *b = reinterpret_cast(b_v); + const auto *dist_table = reinterpret_cast(dist_table_v); + float sum = 0.0f; + for (size_t m = 0; m < num_chunk; ++m) { + size_t idx = m * kTablePerSub + static_cast(a[m]) * kNumCentroids + + static_cast(b[m]); + sum += dist_table[idx]; + } + *out = sum; +} + +void pq_adc_int8_batch_distance(const void **candidates_v, const void *lut_v, + size_t num, size_t num_chunk, float *out) { + constexpr size_t kNumCentroids = 256; + const auto *lut = reinterpret_cast(lut_v); + // candidates_v is const void**, but we need const uint8_t** + // Use an intermediate cast through const char** to avoid aliasing issues. + auto candidates = reinterpret_cast(candidates_v); + + size_t i = 0; + // Main loop: process 4 candidates per iteration. + // Shared LUT base pointer (tab) is computed once per subquantizer, + // reducing redundant pointer arithmetic across the 4 candidates. + for (; i + 4 <= num; i += 4) { + const uint8_t *c0 = candidates[i]; + const uint8_t *c1 = candidates[i + 1]; + const uint8_t *c2 = candidates[i + 2]; + const uint8_t *c3 = candidates[i + 3]; + float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f, d3 = 0.0f; + for (size_t m = 0; m < num_chunk; ++m) { + const float *tab = lut + m * kNumCentroids; + d0 += tab[c0[m]]; + d1 += tab[c1[m]]; + d2 += tab[c2[m]]; + d3 += tab[c3[m]]; + } + out[i] = d0; + out[i + 1] = d1; + out[i + 2] = d2; + out[i + 3] = d3; + } + // Scalar leftover: remaining candidates processed one at a time. + for (; i < num; ++i) { + const uint8_t *code = candidates[i]; + float d = 0.0f; + for (size_t m = 0; m < num_chunk; ++m) { + d += lut[m * kNumCentroids + code[m]]; + } + out[i] = d; + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h new file mode 100644 index 000000000..580361e76 --- /dev/null +++ b/src/turbo/distance/scalar/pq_quantizer_int8/pq_distance.h @@ -0,0 +1,49 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include + +namespace zvec::turbo::scalar { + +// ADC (Asymmetric Distance Computation): compute the distance between a +// PQ-encoded datapoint and a query using a precomputed LUT. +// +// distance = sum_{m=0}^{num_chunk-1} lut[m * 256 + pq_code[m]] +void pq_adc_int8_distance(const void *pq_code, const void *lut, + size_t num_chunk, float *out); + +// SDC (Symmetric Distance Computation): compute the distance between two +// PQ-encoded datapoints using a precomputed centroid-to-centroid distance +// table. +// +// dist_table layout: [num_chunk * 256 * 256] +// dist_table[m * 256 * 256 + i * 256 + j] = +// ||centroid[m][i] - centroid[m][j]||^2 +// +// distance = sum_{m=0}^{num_chunk-1} +// dist_table[m * 65536 + a[m] * 256 + b[m]] +void pq_sdc_int8_distance(const void *a, const void *b, const void *dist_table, + size_t num_chunk, float *out); + +// Batch ADC: compute distances for multiple PQ codes against a shared LUT. +// Processes 4 candidates per iteration (batch4) with shared LUT pointer +// offsets and 4 independent accumulators for ILP. +// Falls back to scalar per-code loop for the remaining candidates. +void pq_adc_int8_batch_distance(const void **candidates, const void *lut, + size_t num, size_t num_chunk, float *out); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc new file mode 100644 index 000000000..d0a5a191b --- /dev/null +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -0,0 +1,855 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "quantizer/pq_int8_quantizer/pq_int8_quantizer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace turbo { + +// --------------------------------------------------------------------------- +// PQ serialization payload (follows the QuantizerSerHeader). +// --------------------------------------------------------------------------- +struct PqInt8SerPayload { + uint32_t original_dim; + uint32_t num_chunk; + uint32_t sub_dim; + uint32_t num_centroids; // always 256 for int8 + uint8_t use_zero_mean; + uint8_t input_data_type; // turbo DataType: kFp32=3, kFp16=2 + uint8_t reserved[2]; +}; + +int PqInt8Quantizer::init(const IndexMeta &meta, const ailego::Params ¶ms) { + meta_ = meta; + + // Map core IndexMeta::DataType to turbo DataType. + if (meta.data_type() == IndexMeta::DataType::DT_FP16) { + input_data_type_ = DataType::kFp16; + } else if (meta.data_type() == IndexMeta::DataType::DT_FP32) { + input_data_type_ = DataType::kFp32; + } else { + return kErrUnsupported; + } + + uint32_t d = meta.dimension(); + original_dim_ = d; + + // Read num_chunk from params (required). + uint32_t nsq = 0; + if (!params.get("num_chunk", &nsq) || nsq == 0) { + return kErrUnsupported; + } + if (d % nsq != 0) { + return kErrUnsupported; + } + + num_chunk_ = nsq; + sub_dim_ = d / nsq; + + // Pre-allocate centroids as raw bytes in the original data type. + centroids_.resize(static_cast(num_chunk_) * kNumCentroids * sub_dim_ * + element_size()); + + // Dispatch ISA kernels (scalar only for now). + auto pq_k = get_pq_kernels(DataType::kInt8); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; + + // Resolve the configured metric type (local only -- not stored). + auto mt = metric_from_name(meta_.metric_name()); + + // L2-only batch distance for encoding: the PQ codebook is trained in L2 + // space regardless of the search metric. Data type matches input. + l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, + QuantizeType::kDefault, CpuArchType::kAuto); + + // Cosine = normalize + L2: after normalization cosine distance is monotonic + // with squared-Euclidean, so the search LUT reuses SquaredEuclidean. + if (meta_.metric_name() == "Cosine") { + batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, + QuantizeType::kDefault, CpuArchType::kAuto); + extra_meta_size_ = kExtraMetaSizeCosine; + meta_.set_extra_meta_size(extra_meta_size_); + } else { + batch_fn_ = get_batch_distance_func( + mt, input_data_type_, QuantizeType::kDefault, CpuArchType::kAuto); + } + + // Read optional training params (aligned with multi_chunk_cluster) + params.get("thread_count", &thread_count_); + params.get("markov_chain_length", &markov_chain_length_); + params.get("epsilon", &epsilon_); + params.get("use_zero_mean", &use_zero_mean_); + + if (use_zero_mean_ && meta_.metric_name() != "SquaredEuclidean" && + meta_.metric_name() != "Cosine") { + use_zero_mean_ = false; + } + + // Set output meta: the quantized representation is INT8 codes with + // num_chunk_ bytes (+ extra_meta_size_ for Cosine norm storage). + meta_.set_meta(IndexMeta::DataType::DT_INT8, num_chunk_); + meta_.set_extra_meta_size(extra_meta_size_); + return 0; +} + +// --------------------------------------------------------------------------- +// Simple Lloyd's KMeans for one sub-quantizer (templated on data type T). +// --------------------------------------------------------------------------- + +template +void PqInt8Quantizer::train_subquantizer(const T *data, size_t num, + size_t stride, size_t sub_idx) { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + uint8_t *centroids_m = + centroids_.data() + static_cast(sub_idx) * k * d * sizeof(T); + + // Non-spherical L2 KMeans: the PQ codebook must minimize L2 reconstruction + // error, so centroids are the true (magnitude-preserving) means. + ailego::NumericalKmeans algorithm(k, d); + + // Append sub-vectors (NumericalKmeans handles transpose internally) + for (size_t i = 0; i < num; ++i) { + const T *sub_vec = + reinterpret_cast(reinterpret_cast(data) + + i * stride) + + sub_idx * d; + algorithm.append(sub_vec, d); + } + + // Single-threaded pool -- parallelism is at the sub-quantizer level. + auto local_threads = std::make_shared(1, false); + + // KMC2 centroid initialization. + ailego::Kmc2CentroidsGenerator< + ailego::NumericalKmeans, + SingleQueueIndexThreads> + gen; + gen.set_chain_length(markov_chain_length_); + gen.set_assumption_free(false); + algorithm.init_centroids(*local_threads, gen); + + // Lloyd iterations + double cost = 0.0; + for (uint32_t iter = 0; iter < kMaxKmeansIters; ++iter) { + double old_cost = cost; + bool result = algorithm.cluster_once(*local_threads, &cost); + if (!result) { + break; + } + double new_epsilon = std::abs(cost - old_cost); + if (new_epsilon < epsilon_) { + break; + } + } + + // Extract centroids into the flat centroids_ byte buffer + const auto ¢s = algorithm.centroids(); + for (size_t c = 0; c < cents.count(); ++c) { + std::memcpy(centroids_m + c * d * sizeof(T), cents[c], d * sizeof(T)); + } +} + +int PqInt8Quantizer::train(IndexHolder::Pointer holder) { + return train(holder, static_cast(thread_count_)); +} + +int PqInt8Quantizer::train(IndexHolder::Pointer holder, int thread_count) { + if (!holder) { + return kErrUnsupported; + } + + size_t num = holder->count(); + const uint32_t elem_sz = element_size(); + + // Collect all data into a contiguous byte buffer (original data type). + auto iter = holder->create_iterator(); + std::vector all_data(num * original_dim_ * elem_sz); + size_t row = 0; + for (; iter->is_valid(); iter->next(), ++row) { + std::memcpy(all_data.data() + row * original_dim_ * elem_sz, iter->data(), + original_dim_ * elem_sz); + } + + // Subsample if the dataset exceeds the training limit (aligned with + // faiss/vsag: 256 centroids * 256 max_points_per_centroid ~= 65535). + if (num > kMaxTrainVectors) { + std::mt19937 rng(42); + // Fisher-Yates partial shuffle: randomly place kMaxTrainVectors vectors + // at the front of the buffer. + for (size_t i = 0; i < kMaxTrainVectors; ++i) { + std::uniform_int_distribution dist(i, num - 1); + size_t j = dist(rng); + if (i != j) { + // Swap full vectors (dim-sized chunks in bytes). + size_t vec_bytes = original_dim_ * elem_sz; + for (size_t b = 0; b < vec_bytes; ++b) { + std::swap(all_data[i * vec_bytes + b], all_data[j * vec_bytes + b]); + } + } + } + num = kMaxTrainVectors; + all_data.resize(num * original_dim_ * elem_sz); + all_data.shrink_to_fit(); + } + + size_t data_stride = original_dim_ * elem_sz; + + // For Cosine: normalize training data so centroids are learned in + // normalized space (L2 minimization == maximizing cosine similarity). + if (meta_.metric_name() == "Cosine") { + switch (input_data_type_) { + case DataType::kFp16: + normalize_batch(reinterpret_cast(all_data.data()), + num); + break; + case DataType::kFp32: + normalize_batch(reinterpret_cast(all_data.data()), num); + break; + default: + break; + } + } + + // Zero-mean centering: subtract the per-dimension mean; the centroid is + // saved for quantize_data / quantize_query / dequantize. For Cosine this + // runs AFTER normalization, so all paths keep the same order + // (normalize -> center; dequantize: un-center -> rescale). + if (use_zero_mean_) { + switch (input_data_type_) { + case DataType::kFp16: + compute_and_subtract_center( + reinterpret_cast(all_data.data()), num); + break; + case DataType::kFp32: + compute_and_subtract_center(reinterpret_cast(all_data.data()), + num); + break; + default: + break; + } + } + + // Create thread pool. + auto threads = std::make_shared( + static_cast(thread_count), false); + auto task_group = threads->make_group(); + + // Distribute sub-quantizers across threads, dispatched by data type. + std::atomic finished{0}; + size_t pool_count = threads->count(); + + auto submit_training = [&](const auto *typed_data) { + using T = std::remove_const_t>; + for (size_t i = 0; i < pool_count; ++i) { + task_group->submit(ailego::Closure::New( + [this, typed_data, num, data_stride, i, pool_count, &finished]() { + for (uint32_t m = static_cast(i); m < num_chunk_; + m += static_cast(pool_count)) { + train_subquantizer(typed_data, num, data_stride, m); + finished++; + } + })); + } + }; + + switch (input_data_type_) { + case DataType::kFp16: + submit_training( + reinterpret_cast(all_data.data())); + break; + case DataType::kFp32: + submit_training(reinterpret_cast(all_data.data())); + break; + default: + break; + } + + task_group->wait_finish(); + + // Pre-build centroid pointer cache (needed by compute_dist_table). + build_centroid_ptrs_cache(); + + // Pre-compute SDC dist_table. + compute_dist_table(); + + return 0; +} + +void PqInt8Quantizer::build_centroid_ptrs_cache() { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + const size_t type_sz = element_size(); + const uint8_t *base = centroids_.data(); + + centroid_ptrs_cache_.resize(num_chunk_); + for (uint32_t m = 0; m < num_chunk_; ++m) { + auto &ptrs = centroid_ptrs_cache_[m]; + ptrs.resize(k); + for (size_t c = 0; c < k; ++c) { + ptrs[c] = base + (static_cast(m) * k * d + c * d) * type_sz; + } + } +} + +void PqInt8Quantizer::compute_dist_table() { + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + dist_table_.resize(static_cast(num_chunk_) * k * k, 0.0f); + + // Centroid-to-centroid distances via the metric-aware batch_fn_: + // L2: dist_table[m][i][j] = ||c_m[i] - c_m[j]||^2 + // IP: dist_table[m][i][j] = -dot(c_m[i], c_m[j]) + // Cosine: centroids trained on normalized data, uses L2. + for (uint32_t m = 0; m < num_chunk_; ++m) { + float *table_m = dist_table_.data() + m * k * k; + + // Use pre-built centroid pointer cache. + // const_cast: .data() returns const void* const* but batch_fn_ + // expects const void**. The kernel never modifies the pointer array. + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + const void *centroid_i = centroid_ptrs[0]; + for (uint32_t i = 0; i < k; ++i) { + batch_fn_(const_cast(centroid_ptrs.data()), + reinterpret_cast(centroid_i) + + static_cast(i) * d * element_size(), + k, d, table_m + i * k); + } + } +} + +void PqInt8Quantizer::quantize_data(const void *input, void *output) const { + uint8_t *code = reinterpret_cast(output); + const uint32_t elem_sz = element_size(); + + // For Cosine: normalize FIRST (codebook is trained in normalized space); + // the original norm is stored after the PQ code for dequantize(). + std::vector norm_vec_storage; + float vec_norm = 0.0f; + const void *vec = input; + + if (meta_.metric_name() == "Cosine") { + norm_vec_storage.resize(original_dim_ * elem_sz); + std::memcpy(norm_vec_storage.data(), input, original_dim_ * elem_sz); + switch (input_data_type_) { + case DataType::kFp16: + normalize_single( + reinterpret_cast(norm_vec_storage.data()), + &vec_norm); + break; + case DataType::kFp32: + normalize_single(reinterpret_cast(norm_vec_storage.data()), + &vec_norm); + break; + default: + break; + } + vec = norm_vec_storage.data(); + } + + // Zero-mean centering: subtract centroid before encoding. + std::vector centered_vec_storage; + if (use_zero_mean_) { + centered_vec_storage.resize(original_dim_ * elem_sz); + std::memcpy(centered_vec_storage.data(), vec, original_dim_ * elem_sz); + switch (input_data_type_) { + case DataType::kFp16: + subtract_center( + reinterpret_cast(centered_vec_storage.data())); + break; + case DataType::kFp32: + subtract_center(reinterpret_cast(centered_vec_storage.data())); + break; + default: + break; + } + vec = centered_vec_storage.data(); + } + + // Encode with L2-only batch distance (search-metric independent), + // fusing argmin into the distance loop. + float dists[kNumCentroids]; + const uint8_t *vec_bytes = reinterpret_cast(vec); + + for (uint32_t m = 0; m < num_chunk_; ++m) { + const void *sub_vec = + vec_bytes + static_cast(m) * sub_dim_ * elem_sz; + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + + // Compute L2 distances from this sub-vector to all 256 centroids. + l2_batch_fn_(const_cast(centroid_ptrs.data()), sub_vec, + kNumCentroids, sub_dim_, dists); + + // Argmin: find nearest centroid. + float best_dist = dists[0]; + uint32_t best_idx = 0; + for (uint32_t j = 1; j < kNumCentroids; ++j) { + if (dists[j] < best_dist) { + best_dist = dists[j]; + best_idx = j; + } + } + code[m] = static_cast(best_idx); + } + + // Store norm after PQ code for Cosine dequantize support. + if (meta_.metric_name() == "Cosine") { + float *norm_out = reinterpret_cast(code + num_chunk_); + *norm_out = vec_norm; + } +} + +void PqInt8Quantizer::quantize_query(const void *input, void *output) const { + float *lut = reinterpret_cast(output); + const uint32_t elem_sz = element_size(); + + // For Cosine: normalize FIRST (Cosine uses an L2 LUT on normalized data), + // consistent with train() / quantize_data(). + std::vector norm_query_storage; + const void *query = input; + + if (meta_.metric_name() == "Cosine") { + norm_query_storage.resize(original_dim_ * elem_sz); + std::memcpy(norm_query_storage.data(), input, original_dim_ * elem_sz); + switch (input_data_type_) { + case DataType::kFp16: + normalize_single( + reinterpret_cast(norm_query_storage.data())); + break; + case DataType::kFp32: + normalize_single(reinterpret_cast(norm_query_storage.data())); + break; + default: + break; + } + query = norm_query_storage.data(); + } + + // Zero-mean centering: subtract centroid before LUT computation. + std::vector centered_query_storage; + if (use_zero_mean_) { + centered_query_storage.resize(original_dim_ * elem_sz); + std::memcpy(centered_query_storage.data(), query, original_dim_ * elem_sz); + switch (input_data_type_) { + case DataType::kFp16: + subtract_center( + reinterpret_cast(centered_query_storage.data())); + break; + case DataType::kFp32: + subtract_center( + reinterpret_cast(centered_query_storage.data())); + break; + default: + break; + } + query = centered_query_storage.data(); + } + + // LUT[m][j] = distance(q_m, c_m[j]) via the metric-aware batch_fn_. + // L2/Cosine: ||q_m - c_m[j]||^2 IP: -dot(q_m, c_m[j]). + // const_cast: see compute_dist_table for rationale. + const uint8_t *query_bytes = reinterpret_cast(query); + for (uint32_t m = 0; m < num_chunk_; ++m) { + const auto ¢roid_ptrs = centroid_ptrs_cache_[m]; + const void *sub_query = + query_bytes + static_cast(m) * sub_dim_ * elem_sz; + batch_fn_(const_cast(centroid_ptrs.data()), sub_query, + kNumCentroids, sub_dim_, lut + m * kNumCentroids); + } + + // Cosine: the LUT holds ||q_m - c_m[j]||^2 on L2-normalized vectors, and + // ||q - c||^2 = 2 - 2*cos_sim = 2*(1 - cos_sim). Scale by 0.5 so the ADC + // sum yields cosine distance (1 - cos_sim) directly on every ADC path. + if (meta_.metric_name() == "Cosine") { + const size_t lut_size = static_cast(num_chunk_) * kNumCentroids; + for (size_t i = 0; i < lut_size; ++i) { + lut[i] *= 0.5f; + } + } +} + +float PqInt8Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + // dp = pq_code (uint8_t[num_chunk]) + // query = LUT (float[num_chunk * 256]) + float d = 0.0f; + adc_fn_(reinterpret_cast(dp), + reinterpret_cast(query), num_chunk_, &d); + // Cosine LUT is pre-scaled by 0.5 in quantize_query, so the ADC sum is + // already the cosine distance -- no conversion needed here. + return d; +} + +void PqInt8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, + const void *query, + float *dist_list) const { + // ISA-dispatched batch4 ADC kernel (4-way ILP + SIMD gather). + // const_cast: batch_adc_fn_ expects const void**; kernel is read-only. + batch_adc_fn_(const_cast(dp_list), query, + static_cast(dp_num), num_chunk_, dist_list); + // Cosine LUT is pre-scaled by 0.5 in quantize_query, so the batch ADC sums + // are already cosine distances -- no conversion applied here. +} + +float PqInt8Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + // Build LUT on the fly, then use ADC. + std::vector lut(static_cast(num_chunk_) * kNumCentroids); + quantize_query(query, lut.data()); + float d = 0.0f; + adc_fn_(reinterpret_cast(dp), lut.data(), num_chunk_, &d); + return d; +} + +void PqInt8Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::vector lut(static_cast(num_chunk_) * kNumCentroids); + quantize_query(query, lut.data()); + // Use ISA-dispatched batch4 ADC kernel (4-way ILP + SIMD gather). + // const_cast: see calc_distance_dp_query_batch for rationale. + batch_adc_fn_(const_cast(dp_list), lut.data(), + static_cast(dp_num), num_chunk_, dist_list); +} + +float PqInt8Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + float d = 0.0f; + sdc_fn_(reinterpret_cast(dp1), + reinterpret_cast(dp2), dist_table_.data(), + num_chunk_, &d); + return d; +} + +int PqInt8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + // Validate unit_size against the input data type. + size_t expected_unit = 0; + switch (input_data_type_) { + case DataType::kFp16: + expected_unit = sizeof(ailego::Float16); + break; + case DataType::kFp32: + expected_unit = sizeof(float); + break; + default: + break; + } + if (qmeta.unit_size() != expected_unit) { + return kErrUnsupported; + } + + size_t lut_bytes = quantized_query_vector_length(); + out->resize(lut_bytes); + quantize_query(query, &(*out)[0]); + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_FP32, original_dim_, + static_cast(type_), 0); + return 0; +} + +int PqInt8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + (void)qmeta; + const uint8_t *code = reinterpret_cast(in); + size_t byte_size = static_cast(original_dim_) * sizeof(float); + out->resize(byte_size); + float *result = reinterpret_cast(&(*out)[0]); + + // Reconstruct by concatenating the selected centroids per sub-quantizer, + // converting from the original data type to float. + const size_t k = kNumCentroids; + const size_t d = sub_dim_; + const uint32_t elem_sz = element_size(); + + for (uint32_t m = 0; m < num_chunk_; ++m) { + const uint8_t *centroids_m = + centroids_.data() + static_cast(m) * k * d * elem_sz; + const uint8_t *centroid = + centroids_m + static_cast(code[m]) * d * elem_sz; + + switch (input_data_type_) { + case DataType::kFp16: { + const ailego::Float16 *src = + reinterpret_cast(centroid); + for (size_t j = 0; j < d; ++j) { + result[m * d + j] = static_cast(src[j]); + } + break; + } + case DataType::kFp32: + std::memcpy(result + m * d, centroid, d * sizeof(float)); + break; + default: + break; + } + } + + // Undo zero-mean centering: add the centroid back FIRST (centering was + // applied in the original-type space during encode). + if (use_zero_mean_) { + for (uint32_t j = 0; j < original_dim_; ++j) { + result[j] += centroid_[j]; + } + } + + // For Cosine: rescale the un-centered unit-space vector back to the + // original magnitude using the stored norm. + if (meta_.metric_name() == "Cosine") { + float norm = 0.0f; + std::memcpy(&norm, code + num_chunk_, sizeof(float)); + for (uint32_t j = 0; j < original_dim_; ++j) { + result[j] *= norm; + } + } + return 0; +} + +DistanceImpl PqInt8Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + (void)qmeta; + + // ADC: PqAdcDistanceFunc matches DistanceFunc directly (no lambda needed). + DistanceFunc adc_func = adc_fn_; + + // Batch ADC: ISA-dispatched batch4 kernel, no lambda needed. + BatchDistanceFunc batch_func = batch_adc_fn_; + + // The query is already quantized (LUT) by the caller; copy it directly. + size_t lut_bytes = quantized_query_vector_length(); + std::string lut_storage(static_cast(query), lut_bytes); + + return DistanceImpl(std::move(adc_func), std::move(batch_func), + std::move(lut_storage), static_cast(num_chunk_)); +} + +DistanceImpl PqInt8Quantizer::sym_distance(const void *query, + const IndexQueryMeta &qmeta) const { + (void)qmeta; + + // SDC kernel needs a lambda: 5 parameters (extra dist_table pointer) + // vs DistanceFunc's 4. + auto sdc = sdc_fn_; + const void *dt = dist_table_.data(); + DistanceFunc sdc_func = [sdc, dt](const void *a, const void *b, size_t dim, + float *out) { sdc(a, b, dt, dim, out); }; + + // The query is a PQ code (uint8_t[num_chunk]), NOT a LUT: SDC compares two + // PQ codes via the centroid-to-centroid dist_table_. + size_t code_bytes = quantized_datapoint_vector_length(); + std::string code_storage(static_cast(query), code_bytes); + + // SDC has no batch kernel -- use the 3-arg constructor. + return DistanceImpl(std::move(sdc_func), std::move(code_storage), + static_cast(num_chunk_)); +} + +// --------------------------------------------------------------------------- +// Serialization +// --------------------------------------------------------------------------- +int PqInt8Quantizer::serialize(std::string *out) const { + if (!out) return kErrUnsupported; + + QuantizerSerHeader hdr{}; + hdr.magic = kQuantizerMagic; + hdr.version = kQuantizerSerVersion; + hdr.quant_type = static_cast(QuantizeType::kPQ); + hdr.dim = original_dim_; + hdr.metric = static_cast(metric_from_name(meta_.metric_name())); + + PqInt8SerPayload payload{}; + payload.original_dim = original_dim_; + payload.num_chunk = num_chunk_; + payload.sub_dim = sub_dim_; + payload.num_centroids = kNumCentroids; + payload.use_zero_mean = use_zero_mean_ ? 1 : 0; + payload.input_data_type = static_cast(input_data_type_); + + size_t centroids_bytes = centroids_.size(); + size_t centroid_bytes = use_zero_mean_ ? centroid_.size() * sizeof(float) : 0; + hdr.payload_size = + static_cast(sizeof(payload) + centroids_bytes + centroid_bytes); + + out->clear(); + out->append(reinterpret_cast(&hdr), sizeof(hdr)); + out->append(reinterpret_cast(&payload), sizeof(payload)); + out->append(reinterpret_cast(centroids_.data()), + centroids_bytes); + // Append zero-mean centroid for centering support. + if (use_zero_mean_) { + out->append(reinterpret_cast(centroid_.data()), + centroid_bytes); + } + // dist_table_ is NOT serialized: it is a build-phase-only derivative of the + // codebook (used by SDC), recomputable on demand and unneeded after + // deserialization (search uses ADC). + return 0; +} + +int PqInt8Quantizer::deserialize(std::string &in) { + return deserialize(in.data(), in.size()); +} + +int PqInt8Quantizer::deserialize(const void *data, size_t len) { + if (len < sizeof(QuantizerSerHeader) + sizeof(PqInt8SerPayload)) { + return kErrUnsupported; + } + + const char *ptr = reinterpret_cast(data); + QuantizerSerHeader hdr; + std::memcpy(&hdr, ptr, sizeof(hdr)); + ptr += sizeof(hdr); + + if (hdr.magic != kQuantizerMagic) return kErrUnsupported; + + PqInt8SerPayload payload; + std::memcpy(&payload, ptr, sizeof(payload)); + ptr += sizeof(payload); + + original_dim_ = payload.original_dim; + num_chunk_ = payload.num_chunk; + sub_dim_ = payload.sub_dim; + + // Restore input data type. Old payloads have input_data_type == 0 + // (was reserved), which maps to kInt4 -- treat as kFp32 for compat. + if (payload.input_data_type == 0 || + payload.input_data_type == static_cast(DataType::kInt4) || + payload.input_data_type == static_cast(DataType::kInt8)) { + input_data_type_ = DataType::kFp32; + } else { + input_data_type_ = static_cast(payload.input_data_type); + } + + // Restore centroids (raw bytes in original data type). + size_t centroids_bytes = static_cast(num_chunk_) * kNumCentroids * + sub_dim_ * element_size(); + centroids_.resize(centroids_bytes); + std::memcpy(centroids_.data(), ptr, centroids_bytes); + ptr += centroids_bytes; + + // Restore zero-mean centroid if centering was enabled. + if (payload.use_zero_mean) { + use_zero_mean_ = true; + size_t centroid_bytes = static_cast(original_dim_) * sizeof(float); + centroid_.resize(original_dim_); + std::memcpy(centroid_.data(), ptr, centroid_bytes); + ptr += centroid_bytes; + } + // dist_table_ is intentionally not restored: SDC is only needed during + // offline build, not after deserialization (search uses ADC). + + // Re-dispatch kernels. + auto pq_k = get_pq_kernels(DataType::kInt8); + adc_fn_ = pq_k.adc_distance; + sdc_fn_ = pq_k.sdc_distance; + batch_adc_fn_ = pq_k.batch_adc_distance; + + // L2-only batch distance for encoding (always L2 regardless of metric). + l2_batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, + QuantizeType::kDefault, CpuArchType::kAuto); + + // Metric-aware batch distance for search LUT. Cosine = normalize + L2, + // so it uses SquaredEuclidean (same as encoding), not IP. + if (meta_.metric_name() == "Cosine") { + batch_fn_ = + get_batch_distance_func(MetricType::kSquaredEuclidean, input_data_type_, + QuantizeType::kDefault, CpuArchType::kAuto); + extra_meta_size_ = kExtraMetaSizeCosine; + } else { + batch_fn_ = get_batch_distance_func( + metric_from_name(meta_.metric_name()), input_data_type_, + QuantizeType::kDefault, CpuArchType::kAuto); + } + + // Set output meta: the quantized representation is INT8 codes with + // num_chunk_ bytes (+ extra_meta_size_ for Cosine norm storage). + meta_.set_meta(IndexMeta::DataType::DT_INT8, num_chunk_); + meta_.set_extra_meta_size(extra_meta_size_); + + // Pre-build centroid pointer cache for fast encode/search. + build_centroid_ptrs_cache(); + + return 0; +} + +INDEX_FACTORY_REGISTER_QUANTIZER(PqInt8Quantizer); + +// --------------------------------------------------------------------------- +// Template helper implementations (type-dispatched at call sites) +// --------------------------------------------------------------------------- + +template +void PqInt8Quantizer::normalize_batch(T *data, size_t num) const { + for (size_t i = 0; i < num; ++i) { + float norm = 0.0f; + ailego::Normalizer::L2(data + i * original_dim_, original_dim_, &norm); + } +} + +template +void PqInt8Quantizer::compute_and_subtract_center(T *data, size_t num) { + centroid_.assign(original_dim_, 0.0f); + for (size_t i = 0; i < num; ++i) { + const T *v = data + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] += static_cast(v[d]); + } + } + for (uint32_t d = 0; d < original_dim_; ++d) { + centroid_[d] /= static_cast(num); + } + for (size_t i = 0; i < num; ++i) { + T *v = data + i * original_dim_; + for (uint32_t d = 0; d < original_dim_; ++d) { + v[d] -= centroid_[d]; + } + } +} + +template +void PqInt8Quantizer::normalize_single(T *vec, float *norm_out) const { + float norm = 0.0f; + ailego::Normalizer::L2(vec, original_dim_, &norm); + if (norm_out) { + *norm_out = norm; + } +} + +template +void PqInt8Quantizer::subtract_center(T *vec) const { + for (uint32_t d = 0; d < original_dim_; ++d) { + vec[d] -= centroid_[d]; + } +} + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h new file mode 100644 index 000000000..0ad673c25 --- /dev/null +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -0,0 +1,221 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +//! Product Quantizer with 8-bit sub-codes (num_bits=8, 256 centroids). +//! +//! Datapoints are encoded as uint8_t[num_chunk] codes. +//! Queries are encoded as a float LUT of size [num_chunk * 256] +//! via quantize_query(). Distance between a PQ code and a +//! query uses ADC (LUT look-up); distance between two PQ codes uses +//! SDC (centroid-to-centroid distance table). +class PqInt8Quantizer : public Quantizer { + public: + PqInt8Quantizer() { + type_ = QuantizeType::kPQ; + } + + ~PqInt8Quantizer() override = default; + + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + const IndexMeta &meta() const override { + return meta_; + } + + DataType input_data_type() const override { + return input_data_type_; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return true; + } + + int train(IndexHolder::Pointer holder) override; + + int train(IndexHolder::Pointer holder, int thread_count) override; + + size_t quantized_datapoint_vector_length() const override { + return num_chunk_ + extra_meta_size_; + } + + size_t quantized_query_vector_length() const override { + return static_cast(num_chunk_) * kNumCentroids * sizeof(float); + } + + void quantize_data(const void *input, void *output) const override; + + void quantize_query(const void *input, void *output) const override; + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const IndexQueryMeta &qmeta, std::string *out, + IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const IndexQueryMeta &qmeta) const override; + + DistanceImpl sym_distance(const void *query, + const IndexQueryMeta &qmeta) const override; + + int serialize(std::string *out) const override; + + int deserialize(std::string &in) override; + + int deserialize(const void *data, size_t len) override; + + private: + //! Train a single sub-quantizer (KMeans, k=256) on the sub-vectors. + //! Templated on the data type T (float or ailego::Float16) so that + //! NumericalKmeans operates natively in the input precision. + //! sub_idx selects which sub-quantizer to train. + template + void train_subquantizer(const T *data, size_t num, size_t stride, + size_t sub_idx); + + //! L2-normalize a batch of vectors (train-time use). + template + void normalize_batch(T *data, size_t num) const; + + //! Compute the per-dimension mean (accumulated in float to avoid FP16 + //! overflow) and subtract it from all training vectors in-place. + template + void compute_and_subtract_center(T *data, size_t num); + + //! L2-normalize a single vector; optionally writes the norm out. + template + void normalize_single(T *vec, float *norm_out = nullptr) const; + + //! Subtract the pre-computed centroid_ from a single vector. + template + void subtract_center(T *vec) const; + + //! Compute the centroid-to-centroid distance table for SDC. + void compute_dist_table(); + + //! Build centroid_ptrs_cache_ from current centroids_. + //! Called after train() and deserialize() when centroids are available. + void build_centroid_ptrs_cache(); + + //! Byte size of one element in the original data type. + uint32_t element_size() const { + return (input_data_type_ == DataType::kFp16) + ? static_cast(sizeof(ailego::Float16)) + : static_cast(sizeof(float)); + } + + static constexpr uint32_t kNumCentroids = 256; + static constexpr uint32_t kMaxKmeansIters = 25; + static constexpr size_t kMaxTrainVectors = 65536; + static constexpr uint32_t kExtraMetaSizeCosine = sizeof(float); + + //! Actual input data type (kFp32 or kFp16). + DataType input_data_type_{DataType::kFp32}; + + //! Thread count for KMeans training (0 = hardware_concurrency). + //! Read from params in init(), aligned with multi_chunk_cluster. + uint32_t thread_count_{0}; + + //! KMC2 Markov chain length (aligned with multi_chunk_cluster default 32). + uint32_t markov_chain_length_{32}; + + //! Cost-based convergence threshold (aligned with multi_chunk_cluster). + double epsilon_{std::numeric_limits::epsilon()}; + + //! Whether to apply zero-mean centering before training/encoding. + //! When enabled, the per-dimension mean of training data is subtracted + //! from all vectors (train, encode, query) and added back on dequantize. + bool use_zero_mean_{false}; + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + uint32_t num_chunk_{0}; + uint32_t sub_dim_{0}; + + //! Centroids stored as raw bytes in the original data type: + //! [num_chunk * kNumCentroids * sub_dim * sizeof(T)] + //! T = float for kFp32, ailego::Float16 for kFp16. + std::vector centroids_; + + //! Global centroid (per-dimension mean) for zero-mean centering. + //! Size: original_dim_ floats. Only populated when use_zero_mean_ = true. + std::vector centroid_; + + //! Centroid-to-centroid distance table for SDC: + //! [num_chunk * kNumCentroids * kNumCentroids] + std::vector dist_table_; + + //! Pre-built centroid pointer arrays for each sub-quantizer. + //! Layout: centroid_ptrs_cache_[sub_idx][centroid_idx] = pointer to centroid. + //! Built once during init/deserialize, reused by compute_dist_table + //! and quantize_query to avoid repeated allocations. + std::vector> centroid_ptrs_cache_; + + //! ISA-dispatched kernel function pointers (ADC / SDC / Batch ADC). + PqAdcDistanceFunc adc_fn_{nullptr}; + PqSdcKernelFunc sdc_fn_{nullptr}; + PqBatchAdcFunc batch_adc_fn_{nullptr}; + + //! Metric-aware batch distance function for search-side LUT + //! computation and SDC dist_table. Data type matches input_data_type_. + //! Obtained from get_batch_distance_func() with the configured metric. + BatchDistanceFunc batch_fn_{}; + + //! L2-only batch distance function for encoding (quantize_data). + //! Data type matches input_data_type_. PQ encoding always minimizes L2 + //! quantization error regardless of the search metric. + BatchDistanceFunc l2_batch_fn_{}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 77572a6a8..90d866585 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -80,6 +80,12 @@ class Quantizer { return 0; } + //! Train the quantizer with data from an IndexHolder using multiple threads. + //! Default implementation ignores thread_count and delegates to train(). + virtual int train(IndexHolder::Pointer holder, int /*thread_count*/) { + return train(holder); + } + //! Byte length of a quantized datapoint vector virtual size_t quantized_datapoint_vector_length() const = 0; @@ -131,6 +137,11 @@ class Quantizer { return DistanceImpl{}; } + virtual DistanceImpl sym_distance(const void *query, + const IndexQueryMeta &qmeta) const { + return distance(query, qmeta); + } + //! Serialize quantizer parameters virtual int serialize(std::string * /*out*/) const { return 0; diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 5a6f33d49..5ddb55b13 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -15,16 +15,23 @@ #include #include #include +#include "avx2/pq_quantizer_int8/pq_distance.h" #include "avx2/rotate/fht/fht.h" +#include "avx512/pq_quantizer_int8/pq_distance.h" #include "avx512/rotate/fht/fht.h" #include "avx512_vnni/record_quantized_int8/cosine.h" #include "avx512_vnni/record_quantized_int8/squared_euclidean.h" #include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/squared_euclidean.h" +#include "neon/pq_quantizer_int8/pq_distance.h" #include "neon/rotate/fht/fht.h" +#include "scalar/fp16/cosine.h" +#include "scalar/fp16/inner_product.h" +#include "scalar/fp16/squared_euclidean.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" +#include "scalar/pq_quantizer_int8/pq_distance.h" #include "scalar/rotate/fht/fht.h" #include "sse/rotate/fht/fht.h" @@ -53,6 +60,21 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, } return nullptr; } + if (data_type == DataType::kFp16) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp16) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp16_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp16_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp16_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && @@ -96,6 +118,21 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, } return nullptr; } + if (data_type == DataType::kFp16) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kFp16) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp16_batch_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp16_batch_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_fp16_batch_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { if (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && @@ -153,6 +190,32 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } +PqKernels get_pq_kernels(DataType data_type, QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + (void)data_type; + if (quantize_type == QuantizeType::kPQ) { + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { + return {avx512::pq_adc_int8_distance_avx512, + avx512::pq_sdc_int8_distance_avx512, + avx512::pq_adc_int8_batch_distance_avx512}; + } + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { + return {avx2::pq_adc_int8_distance_avx2, avx2::pq_sdc_int8_distance_avx2, + avx2::pq_adc_int8_batch_distance_avx2}; + } + if (zvec::ailego::internal::CpuFeatures::static_flags_.NEON && + IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { + return {neon::pq_adc_int8_distance_neon, neon::pq_sdc_int8_distance_neon, + neon::pq_adc_int8_batch_distance_neon}; + } + return {scalar::pq_adc_int8_distance, scalar::pq_sdc_int8_distance, + scalar::pq_adc_int8_batch_distance}; + } + return {}; +} + RotatorKernels get_rotator_kernels(RotateType rotate_type, CpuArchType cpu_arch_type) { switch (rotate_type) { @@ -177,9 +240,6 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, } } - // Unsupported RotateType: assert in debug for early detection, but always - // return the scalar kernels so release builds never hand back null function - // pointers (which would crash on the first call). assert(false && "unsupported RotateType"); return {scalar::fht_rotate, scalar::fht_unrotate}; } diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 26f199c49..230655cc0 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -2525,7 +2525,88 @@ TEST(IndexInterface, BuilderSetsAllBaseFields) { EXPECT_FALSE(param->use_id_map); EXPECT_TRUE(param->use_external_vector); EXPECT_EQ(PreprocessorType::kPCA, param->preprocess_param.type); - EXPECT_EQ(QuantizerType::kFP16, param->quantizer_param.type); + EXPECT_EQ(QuantizerType::kFP16, param->quantizer_param->type); +} + +TEST(IndexInterface, QuantizerParamDefaultIsNull) { + auto param = FlatIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDimension(64) + .Build(); + + ASSERT_NE(nullptr, param); + EXPECT_EQ(nullptr, param->quantizer_param); + EXPECT_EQ(QuantizerType::kNone, param->quantizer_type()); + EXPECT_FALSE(param->enable_rotate()); +} + +TEST(IndexInterface, QuantizerParamPqFields) { + auto param = FlatIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDimension(128) + .WithQuantizerParam(PqQuantizerParam(16, 8)) + .Build(); + + ASSERT_NE(nullptr, param); + ASSERT_NE(nullptr, param->quantizer_param); + EXPECT_EQ(QuantizerType::kPQ, param->quantizer_type()); + + // the builder must keep the concrete type instead of slicing it + auto pq_param = + std::dynamic_pointer_cast(param->quantizer_param); + ASSERT_NE(nullptr, pq_param); + EXPECT_EQ(16, pq_param->num_chunk); + EXPECT_EQ(8, pq_param->num_bits); + + // enable_rotate is a common field, setting it keeps the concrete type + auto rotated_param = FlatIndexParamBuilder() + .WithQuantizerParam(PqQuantizerParam(16, 8)) + .WithEnableRotate(true) + .Build(); + ASSERT_NE(nullptr, rotated_param->quantizer_param); + EXPECT_TRUE(rotated_param->enable_rotate()); + EXPECT_NE(nullptr, std::dynamic_pointer_cast( + rotated_param->quantizer_param)); +} + +TEST(IndexInterface, QuantizerParamPqJsonRoundTrip) { + auto param = FlatIndexParamBuilder() + .WithIndexType(IndexType::kFlat) + .WithMetricType(MetricType::kL2sq) + .WithDimension(128) + .WithDataType(DataType::DT_FP32) + .WithQuantizerParam(PqQuantizerParam(32, 4)) + .Build(); + + auto deserialized_param = + IndexFactory::DeserializeIndexParamFromJson(param->SerializeToJson()); + ASSERT_NE(nullptr, deserialized_param); + EXPECT_EQ(param->SerializeToJson(), deserialized_param->SerializeToJson()); + EXPECT_EQ(param->SerializeToJson(true), + deserialized_param->SerializeToJson(true)); + + auto pq_param = std::dynamic_pointer_cast( + deserialized_param->quantizer_param); + ASSERT_NE(nullptr, pq_param); + EXPECT_EQ(QuantizerType::kPQ, pq_param->type); + EXPECT_EQ(32, pq_param->num_chunk); + EXPECT_EQ(4, pq_param->num_bits); +} + +TEST(IndexInterface, QuantizerParamLegacyJsonCompat) { + // legacy json only carries the common fields + const std::string json_str = + R"({"index_type":"kFlat","metric_type":"kL2sq","dimension":128,)" + R"("data_type":"DT_FP32",)" + R"("quantizer_param":{"type":"kInt8","enable_rotate":true}})"; + + auto param = IndexFactory::DeserializeIndexParamFromJson(json_str); + ASSERT_NE(nullptr, param); + ASSERT_NE(nullptr, param->quantizer_param); + EXPECT_EQ(QuantizerType::kInt8, param->quantizer_type()); + EXPECT_TRUE(param->enable_rotate()); + EXPECT_EQ(nullptr, std::dynamic_pointer_cast( + param->quantizer_param)); } TEST(IndexInterface, BuilderChainingReturnsCorrectType) { diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc new file mode 100644 index 000000000..627e1ec5f --- /dev/null +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -0,0 +1,1449 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include "distance/scalar/pq_quantizer_int8/pq_distance.h" +#include "quantizer/pq_int8_quantizer/pq_int8_quantizer.h" +#include "zvec/core/framework/index_factory.h" + +#if defined(__AVX2__) +#include "distance/avx2/pq_quantizer_int8/pq_distance.h" +#endif +#if defined(__AVX512F__) +#include "distance/avx512/pq_quantizer_int8/pq_distance.h" +#endif + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; +using zvec::turbo::DataType; + +// Reference squared Euclidean distance between two raw fp32 vectors. +static float reference_sq_euclidean(const float *a, const float *b, + size_t dim) { + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + float diff = a[i] - b[i]; + sum += diff * diff; + } + return sum; +} + +// Helper to create a PqInt8Quantizer via the factory. +static std::shared_ptr make_pq_quantizer( + size_t dim, size_t num_chunk) { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("SquaredEuclidean", 0, Params()); + + Params params; + params.set("num_chunk", static_cast(num_chunk)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Helper: build a holder with random fp32 vectors. +static std::shared_ptr> +make_random_holder(size_t count, size_t dim, uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen); + holder->emplace(i + 1, vec); + } + return holder; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +TEST(PqInt8Quantizer, InitInvalidParams) { + // dim not divisible by num_chunk + auto q = make_pq_quantizer(10, 3); + EXPECT_EQ(q, nullptr); + + // num_chunk = 0 + auto q2 = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_TRUE(q2); + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, 16); + meta.set_metric("SquaredEuclidean", 0, Params()); + Params params; + params.set("num_chunk", static_cast(0)); + EXPECT_NE(0, q2->init(meta, params)); +} + +TEST(PqInt8Quantizer, TrainAndEncode) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + EXPECT_TRUE(quantizer->require_train()); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Quantize a few vectors and check code length. + auto iter = holder->create_iterator(); + size_t checked = 0; + std::vector code(quantizer->quantized_datapoint_vector_length()); + for (; iter->is_valid() && checked < 10; iter->next(), ++checked) { + quantizer->quantize_data(iter->data(), code.data()); + // Each code byte should be in [0, 255]. + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_LE(code[m], 255u); + } + } + EXPECT_EQ(10u, checked); +} + +TEST(PqInt8Quantizer, AdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0] + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + // ADC distances should be a reasonable approximation of true distance. + // With 8 sub-quantizers and 32 dims (sub_dim=4), PQ error is non-trivial + // but should be bounded. + float max_rel_error = 0.0f; + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_sq_euclidean(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + if (true_dist > 1e-6f) { + float rel = std::fabs(adc_dist - true_dist) / true_dist; + max_rel_error = std::max(max_rel_error, rel); + } + // ADC distance must be non-negative. + EXPECT_GE(adc_dist, 0.0f) << "i=" << i; + } + // With 8 subs and 2000 training points, max relative error should be + // well below 100% (generous bound; actual error is typically <30%). + EXPECT_LT(max_rel_error, 1.0f) << "max_rel_error=" << max_rel_error; +} + +TEST(PqInt8Quantizer, SdcDistance) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 2000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Encode two vectors and compute SDC distance. + auto iter = holder->create_iterator(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(quantizer->quantized_datapoint_vector_length()); + + iter->is_valid(); + quantizer->quantize_data(iter->data(), code1.data()); + iter->next(); + iter->is_valid(); + quantizer->quantize_data(iter->data(), code2.data()); + + float sdc_dist = quantizer->calc_distance_dp_dp(code1.data(), code2.data()); + EXPECT_GE(sdc_dist, 0.0f); +} + +TEST(PqInt8Quantizer, DistanceImplAdcAndSdc) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Quantize query[0] as LUT. + auto iter = holder->create_iterator(); + iter->is_valid(); + const float *query_raw = reinterpret_cast(iter->data()); + + size_t lut_bytes = quantizer->quantized_query_vector_length(); + std::string lut_storage(lut_bytes, '\0'); + quantizer->quantize_query(query_raw, &lut_storage[0]); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + auto dist_impl = quantizer->distance(lut_storage.data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + + // func() should be set (ADC path). + EXPECT_TRUE(static_cast(dist_impl.func())); + + // Encode a candidate and compute distance via DistanceImpl (ADC path). + iter->next(); + iter->is_valid(); + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + float d = dist_impl(code.data()); + EXPECT_GE(d, 0.0f); +} + +TEST(PqInt8Quantizer, SerializeDeserialize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 500; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Serialize. + std::string blob; + ASSERT_EQ(0, quantizer->serialize(&blob)); + EXPECT_GT(blob.size(), sizeof(zvec::turbo::QuantizerSerHeader)); + + // Deserialize into a fresh quantizer. + auto q2 = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_TRUE(q2); + ASSERT_EQ(0, q2->deserialize(blob)); + + // Encode the same vector with both and compare codes. + auto iter = holder->create_iterator(); + iter->is_valid(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code1.data()); + q2->quantize_data(iter->data(), code2.data()); + + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_EQ(code1[m], code2[m]) << "m=" << m; + } + + // ADC distances should also match (same codebook → same LUT → same ADC). + size_t lut_len = quantizer->quantized_query_vector_length(); + std::vector lut1(lut_len / sizeof(float)); + std::vector lut2(lut_len / sizeof(float)); + quantizer->quantize_query(iter->data(), lut1.data()); + q2->quantize_query(iter->data(), lut2.data()); + + float adc1 = quantizer->calc_distance_dp_query(code1.data(), lut1.data()); + float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); + EXPECT_NEAR(adc1, adc2, 1e-6f); + + // Note: SDC (calc_distance_dp_dp) is intentionally NOT tested after + // deserialization because dist_table_ is a build-phase-only structure + // and is not persisted. +} + +// --------------------------------------------------------------------------- +// SIMD Consistency Tests +// --------------------------------------------------------------------------- + +namespace { + +// Helper to generate random uint8 codes +void fill_random_codes(uint8_t *codes, size_t len, std::mt19937 &gen) { + std::uniform_int_distribution dist(0, 255); + for (size_t i = 0; i < len; ++i) { + codes[i] = static_cast(dist(gen)); + } +} + +// Helper to generate random LUT (ADC) +void fill_random_lut(float *lut, size_t num_chunk, std::mt19937 &gen) { + constexpr size_t kNumCentroids = 256; + std::uniform_real_distribution dist(0.0f, 1.0f); + for (size_t m = 0; m < num_chunk; ++m) { + for (size_t c = 0; c < kNumCentroids; ++c) { + lut[m * kNumCentroids + c] = dist(gen); + } + } +} + +// Helper to generate random dist_table (SDC) +void fill_random_sdc_table(float *table, size_t num_chunk, std::mt19937 &gen) { + constexpr size_t kTablePerSub = 256 * 256; + std::uniform_real_distribution dist(0.0f, 1.0f); + for (size_t m = 0; m < num_chunk; ++m) { + for (size_t i = 0; i < kTablePerSub; ++i) { + table[m * kTablePerSub + i] = dist(gen); + } + } +} + +} // anonymous namespace + +// Test ADC SIMD consistency across multiple M values +TEST(PqInt8SimdConsistency, AdcDistance) { + std::mt19937 gen(2024); + + // Test various M values including boundary cases + // M=4,8: exact multiples of AVX2 chunk (8) + // M=12: not multiple of 8, has leftover + // M=16: exact multiple of AVX512 chunk (16) + for (size_t num_sq : {4, 8, 12, 16}) { + constexpr size_t kNumCentroids = 256; + std::vector codes(num_sq); + std::vector lut(num_sq * kNumCentroids); + + fill_random_codes(codes.data(), num_sq, gen); + fill_random_lut(lut.data(), num_sq, gen); + + // Compute reference (scalar) + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_adc_int8_distance(codes.data(), lut.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_adc_int8_distance_avx2(codes.data(), lut.data(), + num_sq, &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) + << "AVX2 ADC mismatch for M=" << num_sq; + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_adc_int8_distance_avx512(codes.data(), lut.data(), + num_sq, &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) + << "AVX512 ADC mismatch for M=" << num_sq; + } +#endif + } +} + +// Test SDC SIMD consistency across multiple M values +TEST(PqInt8SimdConsistency, SdcDistance) { + std::mt19937 gen(2025); + + for (size_t num_sq : {4, 8, 12, 16}) { + constexpr size_t kTablePerSub = 256 * 256; + std::vector codes_a(num_sq); + std::vector codes_b(num_sq); + std::vector dist_table(num_sq * kTablePerSub); + + fill_random_codes(codes_a.data(), num_sq, gen); + fill_random_codes(codes_b.data(), num_sq, gen); + fill_random_sdc_table(dist_table.data(), num_sq, gen); + + // Compute reference (scalar) + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_sdc_int8_distance(codes_a.data(), codes_b.data(), + dist_table.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_sdc_int8_distance_avx2( + codes_a.data(), codes_b.data(), dist_table.data(), num_sq, + &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f) + << "AVX2 SDC mismatch for M=" << num_sq; + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_sdc_int8_distance_avx512( + codes_a.data(), codes_b.data(), dist_table.data(), num_sq, + &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f) + << "AVX512 SDC mismatch for M=" << num_sq; + } +#endif + } +} + +// Test edge case: M=1 (minimum valid value) +TEST(PqInt8SimdConsistency, AdcDistanceM1) { + std::mt19937 gen(123); + constexpr size_t kNumCentroids = 256; + constexpr size_t num_sq = 1; + + std::vector codes(num_sq); + std::vector lut(num_sq * kNumCentroids); + + fill_random_codes(codes.data(), num_sq, gen); + fill_random_lut(lut.data(), num_sq, gen); + + float scalar_result = 0.0f; + zvec::turbo::scalar::pq_adc_int8_distance(codes.data(), lut.data(), num_sq, + &scalar_result); + +#if defined(__AVX2__) + { + float avx2_result = 0.0f; + zvec::turbo::avx2::pq_adc_int8_distance_avx2(codes.data(), lut.data(), + num_sq, &avx2_result); + EXPECT_NEAR(scalar_result, avx2_result, 1e-5f); + } +#endif + +#if defined(__AVX512F__) + { + float avx512_result = 0.0f; + zvec::turbo::avx512::pq_adc_int8_distance_avx512(codes.data(), lut.data(), + num_sq, &avx512_result); + EXPECT_NEAR(scalar_result, avx512_result, 1e-5f); + } +#endif +} + +// --------------------------------------------------------------------------- +// Cosine Metric Tests +// --------------------------------------------------------------------------- + +// Helper to create a PqInt8Quantizer with Cosine metric. +static std::shared_ptr make_pq_cosine_quantizer( + size_t dim, size_t num_chunk) { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("Cosine", 0, Params()); + + Params params; + params.set("num_chunk", static_cast(num_chunk)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Reference cosine distance: 1 - (a·b) / (||a|| * ||b||). +static float reference_cosine_distance(const float *a, const float *b, + size_t dim) { + float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + norm_a += a[i] * a[i]; + norm_b += b[i] * b[i]; + } + float denom = std::sqrt(norm_a) * std::sqrt(norm_b); + if (denom < 1e-12f) return 1.0f; + return 1.0f - dot / denom; +} + +// Helper: generate random vectors with varying norms (not unit length). +static std::shared_ptr> +make_cosine_holder(size_t count, size_t dim, uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::uniform_real_distribution scale(0.5f, 5.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + float s = scale(gen); + for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen) * s; + holder->emplace(i + 1, vec); + } + return holder; +} + +// Verify that quantize_data stores the correct L2 norm after the PQ code. +TEST(PqInt8Quantizer, CosineNormStorage) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 500; + + auto quantizer = make_pq_cosine_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // extra_meta_size should be sizeof(float) for Cosine. + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), + NSQ + sizeof(float)); + + auto holder = make_cosine_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::vector code(quantizer->quantized_datapoint_vector_length()); + + for (size_t checked = 0; iter->is_valid() && checked < 10; + iter->next(), ++checked) { + const float *v = reinterpret_cast(iter->data()); + + // Compute expected norm. + float expected_norm_sq = 0.0f; + for (size_t j = 0; j < DIM; ++j) expected_norm_sq += v[j] * v[j]; + float expected_norm = std::sqrt(expected_norm_sq); + + quantizer->quantize_data(iter->data(), code.data()); + + // Read stored norm from after the PQ code. + float stored_norm = 0.0f; + std::memcpy(&stored_norm, code.data() + NSQ, sizeof(float)); + + EXPECT_NEAR(stored_norm, expected_norm, expected_norm * 1e-5f) + << "Norm mismatch at vector " << checked; + } +} + +// Verify that dequantize reconstructs a vector with approximately correct +// direction (cosine similarity close to 1) and magnitude. +TEST(PqInt8Quantizer, CosineDequantize) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_cosine_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_cosine_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + + float max_cos_dist = 0.0f; + float max_norm_rel_error = 0.0f; + + for (size_t i = 0; iter->is_valid() && i < 50; iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + + // Compute original norm. + float orig_norm_sq = 0.0f; + for (size_t j = 0; j < DIM; ++j) orig_norm_sq += v[j] * v[j]; + float orig_norm = std::sqrt(orig_norm_sq); + + // Encode. + std::vector code(code_len); + quantizer->quantize_data(iter->data(), code.data()); + + // Decode. + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + + // Check cosine similarity between original and reconstructed. + float cos_dist = reference_cosine_distance(v, recon, DIM); + max_cos_dist = std::max(max_cos_dist, cos_dist); + + // Check norm of reconstructed vector ≈ original norm. + float recon_norm_sq = 0.0f; + for (size_t j = 0; j < DIM; ++j) recon_norm_sq += recon[j] * recon[j]; + float recon_norm = std::sqrt(recon_norm_sq); + + if (orig_norm > 1e-6f) { + float rel_err = std::fabs(recon_norm - orig_norm) / orig_norm; + max_norm_rel_error = std::max(max_norm_rel_error, rel_err); + } + } + + // PQ with 8 subs and 2000 training vectors: cosine distance should be + // small (< 0.3 is generous; typically < 0.1). + EXPECT_LT(max_cos_dist, 0.3f) << "max_cos_dist=" << max_cos_dist; + + // Reconstructed norm should closely match original. PQ centroid + // concatenation in normalized space is not exactly unit-length, so + // the norm carries the centroid approximation error (~10-15% for + // 8 subs with sub_dim=4 and 2000 training vectors). + EXPECT_LT(max_norm_rel_error, 0.2f) + << "max_norm_rel_error=" << max_norm_rel_error; +} + +// Verify Cosine search distances via ADC are consistent with true cosine +// distance and fall in the expected range [0, 2]. +TEST(PqInt8Quantizer, CosineAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_cosine_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_cosine_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0]. + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_cosine_distance(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + + // Cosine distance should be in [0, 2]. + EXPECT_GE(adc_dist, -0.01f) << "i=" << i; + EXPECT_LE(adc_dist, 2.01f) << "i=" << i; + + // PQ approximation: should be roughly correlated (within 0.5). + EXPECT_LT(std::fabs(adc_dist - true_dist), 0.5f) + << "i=" << i << " adc=" << adc_dist << " true=" << true_dist; + } +} + +// Verify that dequantize for L2 metric (no norm storage) still works. +TEST(PqInt8Quantizer, L2Dequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // L2 metric: no extra meta. + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + const float *orig = reinterpret_cast(iter->data()); + + // L2 PQ reconstruction should be a reasonable approximation. + float recon_err = reference_sq_euclidean(orig, recon, DIM); + float orig_norm = reference_sq_euclidean(orig, orig, DIM); + // Relative reconstruction error should be bounded. + if (orig_norm > 1e-6f) { + EXPECT_LT(recon_err / orig_norm, 1.0f) + << "recon_err=" << recon_err << " orig_norm=" << orig_norm; + } +} + +// --------------------------------------------------------------------------- +// InnerProduct Metric Tests +// --------------------------------------------------------------------------- + +// Helper to create a PqInt8Quantizer with InnerProduct metric. +static std::shared_ptr make_pq_ip_quantizer( + size_t dim, size_t num_chunk) { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("InnerProduct", 0, Params()); + + Params params; + params.set("num_chunk", static_cast(num_chunk)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Reference inner-product distance: -dot(a, b). +static float reference_ip_distance(const float *a, const float *b, size_t dim) { + float dot = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + } + return -dot; +} + +// Verify IP metric: no extra meta, ADC distances approximate true IP. +TEST(PqInt8Quantizer, InnerProductAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_ip_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // IP metric should NOT add extra meta (unlike Cosine). + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0]. + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + float max_abs_error = 0.0f; + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_ip_distance(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + + // IP distance can be positive or negative; relative error is + // meaningless near zero crossings. Use absolute error instead. + float abs_err = std::fabs(adc_dist - true_dist); + max_abs_error = std::max(max_abs_error, abs_err); + } + // PQ IP distance: absolute error should be bounded. + // With 8 subs and dim=32, typical max abs error is a few units. + EXPECT_LT(max_abs_error, static_cast(DIM) * 0.5f) + << "max_abs_error=" << max_abs_error; +} + +// Verify IP dequantize works (same as L2: centroid concat, no norm rescale). +TEST(PqInt8Quantizer, InnerProductDequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_ip_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), NSQ); + + auto holder = make_random_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + const float *orig = reinterpret_cast(iter->data()); + + // IP PQ reconstruction: centroid concat, same as L2. + float recon_err = reference_sq_euclidean(orig, recon, DIM); + float orig_norm = reference_sq_euclidean(orig, orig, DIM); + if (orig_norm > 1e-6f) { + EXPECT_LT(recon_err / orig_norm, 1.0f) + << "recon_err=" << recon_err << " orig_norm=" << orig_norm; + } +} + +// --------------------------------------------------------------------------- +// Zero-Mean Centering Tests +// --------------------------------------------------------------------------- + +// Helper to create a PqInt8Quantizer with zero-mean centering enabled. +static std::shared_ptr make_pq_zero_mean_quantizer( + size_t dim, size_t num_chunk) { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("SquaredEuclidean", 0, Params()); + + Params params; + params.set("num_chunk", static_cast(num_chunk)); + params.set("use_zero_mean", true); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Helper: build a holder with random fp32 vectors that have a large offset +// (non-zero mean). This simulates real-world data where centering helps. +static std::shared_ptr> +make_offset_holder(size_t count, size_t dim, float offset = 10.0f, + uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen) + offset; + holder->emplace(i + 1, vec); + } + return holder; +} + +// Verify basic functionality: train, encode, ADC distance with centering. +TEST(PqInt8Quantizer, ZeroMeanTrainAndEncode) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + EXPECT_TRUE(quantizer->require_train()); + + // Use offset data to exercise the centering path meaningfully. + auto holder = make_offset_holder(COUNT, DIM, 10.0f); + ASSERT_EQ(0, quantizer->train(holder)); + + // Quantize a few vectors and check code length. + auto iter = holder->create_iterator(); + size_t checked = 0; + std::vector code(quantizer->quantized_datapoint_vector_length()); + for (; iter->is_valid() && checked < 10; iter->next(), ++checked) { + quantizer->quantize_data(iter->data(), code.data()); + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_LE(code[m], 255u); + } + } + EXPECT_EQ(10u, checked); +} + +// Verify ADC distance accuracy with centering on offset data. +TEST(PqInt8Quantizer, ZeroMeanAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + // Offset data: centering should improve PQ accuracy for high-offset vectors. + auto holder = make_offset_holder(COUNT, DIM, 10.0f); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const float *v = reinterpret_cast(iter->data()); + raw_vecs[i].assign(v, v + DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = raw_vecs[0]. + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(raw_vecs[0].data(), lut.data()); + + // ADC distances should approximate true L2 distances. + float max_rel_error = 0.0f; + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_sq_euclidean(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + if (true_dist > 1e-6f) { + float rel = std::fabs(adc_dist - true_dist) / true_dist; + max_rel_error = std::max(max_rel_error, rel); + } + EXPECT_GE(adc_dist, 0.0f) << "i=" << i; + } + EXPECT_LT(max_rel_error, 1.0f) << "max_rel_error=" << max_rel_error; +} + +// Verify dequantize correctly adds centroid back. +TEST(PqInt8Quantizer, ZeroMeanDequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + const float OFFSET = 10.0f; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_offset_holder(COUNT, DIM, OFFSET); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + const float *orig = reinterpret_cast(iter->data()); + + // Reconstructed vector should be near the original (which has ~10 offset). + float recon_err = reference_sq_euclidean(orig, recon, DIM); + float orig_norm = reference_sq_euclidean(orig, orig, DIM); + if (orig_norm > 1e-6f) { + EXPECT_LT(recon_err / orig_norm, 1.0f) + << "recon_err=" << recon_err << " orig_norm=" << orig_norm; + } + + // The reconstructed values should be around OFFSET (not near zero), + // confirming that the centroid was added back. + float recon_mean = 0.0f; + for (size_t j = 0; j < DIM; ++j) recon_mean += recon[j]; + recon_mean /= DIM; + EXPECT_GT(recon_mean, OFFSET * 0.5f) + << "Reconstructed mean too low; centroid may not be added back"; +} + +// Verify serialize/deserialize preserves the zero-mean centroid. +TEST(PqInt8Quantizer, ZeroMeanSerializeDeserialize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 500; + + auto quantizer = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_offset_holder(COUNT, DIM, 5.0f); + ASSERT_EQ(0, quantizer->train(holder)); + + // Serialize. + std::string blob; + ASSERT_EQ(0, quantizer->serialize(&blob)); + EXPECT_GT(blob.size(), sizeof(zvec::turbo::QuantizerSerHeader)); + + // Deserialize into a fresh quantizer. + auto q2 = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_TRUE(q2); + ASSERT_EQ(0, q2->deserialize(blob)); + + // Encode the same vector with both and compare codes. + auto iter = holder->create_iterator(); + iter->is_valid(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code1.data()); + q2->quantize_data(iter->data(), code2.data()); + + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_EQ(code1[m], code2[m]) << "m=" << m; + } + + // ADC distances should also match. + size_t lut_len = quantizer->quantized_query_vector_length(); + std::vector lut1(lut_len / sizeof(float)); + std::vector lut2(lut_len / sizeof(float)); + quantizer->quantize_query(iter->data(), lut1.data()); + q2->quantize_query(iter->data(), lut2.data()); + + float adc1 = quantizer->calc_distance_dp_query(code1.data(), lut1.data()); + float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); + EXPECT_NEAR(adc1, adc2, 1e-6f); + + // Dequantize from q2 should also produce vectors in the offset range. + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + std::string decoded; + ASSERT_EQ(0, q2->dequantize(code2.data(), qmeta, &decoded)); + const float *recon = reinterpret_cast(decoded.data()); + float recon_mean = 0.0f; + for (size_t j = 0; j < DIM; ++j) recon_mean += recon[j]; + recon_mean /= DIM; + EXPECT_GT(recon_mean, 2.5f) + << "Deserialized quantizer centroid not restored properly"; +} + +// Verify that centering does not significantly degrade PQ accuracy. +// Centering benefits data with skewed distributions or high-mean non-uniform +// spread; for uniform+offset data, the error should remain comparable. +TEST(PqInt8Quantizer, ZeroMeanAccuracyComparable) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + const float OFFSET = 50.0f; + + // Train without centering. + auto q_no_center = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(q_no_center); + auto holder = make_offset_holder(COUNT, DIM, OFFSET); + ASSERT_EQ(0, q_no_center->train(holder)); + + // Train with centering. + auto q_center = make_pq_zero_mean_quantizer(DIM, NSQ); + ASSERT_TRUE(q_center); + ASSERT_EQ(0, q_center->train(holder)); + + // Compute average reconstruction error for both. + auto iter = holder->create_iterator(); + float err_no_center_sum = 0.0f; + float err_center_sum = 0.0f; + size_t checked = 0; + + size_t code_len_no = q_no_center->quantized_datapoint_vector_length(); + size_t code_len_yes = q_center->quantized_datapoint_vector_length(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, DIM); + + for (; iter->is_valid() && checked < 100; iter->next(), ++checked) { + const float *orig = reinterpret_cast(iter->data()); + + // Without centering. + std::vector code1(code_len_no); + q_no_center->quantize_data(iter->data(), code1.data()); + std::string decoded1; + q_no_center->dequantize(code1.data(), qmeta, &decoded1); + err_no_center_sum += reference_sq_euclidean( + orig, reinterpret_cast(decoded1.data()), DIM); + + // With centering. + std::vector code2(code_len_yes); + q_center->quantize_data(iter->data(), code2.data()); + std::string decoded2; + q_center->dequantize(code2.data(), qmeta, &decoded2); + err_center_sum += reference_sq_euclidean( + orig, reinterpret_cast(decoded2.data()), DIM); + } + + float avg_err_no_center = err_no_center_sum / checked; + float avg_err_center = err_center_sum / checked; + + // Centering should not degrade accuracy by more than 2x. + EXPECT_LT(avg_err_center, avg_err_no_center * 2.0f) + << "Centering degraded accuracy too much: center_err=" << avg_err_center + << " no_center_err=" << avg_err_no_center; +} + +// --------------------------------------------------------------------------- +// FP16 Input Tests +// --------------------------------------------------------------------------- + +#include + +// Helper to create a PqInt8Quantizer with FP16 input type. +static std::shared_ptr make_pq_fp16_quantizer( + size_t dim, size_t num_chunk, const char *metric = "SquaredEuclidean") { + auto q = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + if (!q) return nullptr; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP16, dim); + meta.set_metric(metric, 0, Params()); + + Params params; + params.set("num_chunk", static_cast(num_chunk)); + if (q->init(meta, params) != 0) return nullptr; + return q; +} + +// Helper: build a holder with random fp16 vectors (via Float16). +static std::shared_ptr> +make_random_fp16_holder(size_t count, size_t dim, uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + for (size_t j = 0; j < dim; ++j) vec[j] = ailego::Float16(dist(gen)); + holder->emplace(i + 1, vec); + } + return holder; +} + +// Convert a Float16 vector to fp32 for reference distance computation. +static std::vector fp16_to_fp32(const ailego::Float16 *v, size_t dim) { + std::vector out(dim); + for (size_t i = 0; i < dim; ++i) out[i] = static_cast(v[i]); + return out; +} + +// Verify FP16 init succeeds and output meta is correct. +TEST(PqInt8Fp16, InitAndOutputMeta) { + auto q = make_pq_fp16_quantizer(16, 4); + ASSERT_TRUE(q); + + // input_data_type should be kFp16. + EXPECT_EQ(q->input_data_type(), DataType::kFp16); + + // Output meta: data_type = DT_INT8, dimension = num_chunk. + const auto &meta = q->meta(); + EXPECT_EQ(meta.data_type(), IndexMeta::DataType::DT_INT8); + EXPECT_EQ(meta.dimension(), 4u); + // L2 metric: no extra meta. + EXPECT_EQ(meta.extra_meta_size(), 0u); + EXPECT_EQ(meta.element_size(), 4u); + + // Cosine FP16: extra meta for norm storage. + auto q_cos = make_pq_fp16_quantizer(32, 8, "Cosine"); + ASSERT_TRUE(q_cos); + EXPECT_EQ(q_cos->meta().extra_meta_size(), sizeof(float)); + EXPECT_EQ(q_cos->meta().element_size(), 8u + sizeof(float)); +} + +// Verify basic train + encode with FP16 input. +TEST(PqInt8Fp16, TrainAndEncode) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + EXPECT_TRUE(quantizer->require_train()); + + auto holder = make_random_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + size_t checked = 0; + std::vector code(quantizer->quantized_datapoint_vector_length()); + for (; iter->is_valid() && checked < 10; iter->next(), ++checked) { + quantizer->quantize_data(iter->data(), code.data()); + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_LE(code[m], 255u); + } + } + EXPECT_EQ(10u, checked); +} + +// Verify ADC distances with FP16 input are reasonable. +TEST(PqInt8Fp16, AdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors (as fp32 for reference) and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const ailego::Float16 *v = + reinterpret_cast(iter->data()); + raw_vecs[i] = fp16_to_fp32(v, DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = first vector (must use FP16 data from holder). + auto iter2 = holder->create_iterator(); + iter2->is_valid(); + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(iter2->data(), lut.data()); + + float max_rel_error = 0.0f; + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_sq_euclidean(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + if (true_dist > 1e-6f) { + float rel = std::fabs(adc_dist - true_dist) / true_dist; + max_rel_error = std::max(max_rel_error, rel); + } + EXPECT_GE(adc_dist, 0.0f) << "i=" << i; + } + // FP16 has lower precision than FP32, so allow slightly larger error. + EXPECT_LT(max_rel_error, 1.5f) << "max_rel_error=" << max_rel_error; +} + +// Verify dequantize from FP16 PQ produces reasonable fp32 reconstruction. +TEST(PqInt8Fp16, Dequantize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 1000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + iter->is_valid(); + + const ailego::Float16 *orig_fp16 = + reinterpret_cast(iter->data()); + std::vector orig_fp32 = fp16_to_fp32(orig_fp16, DIM); + + std::vector code(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code.data()); + + // Dequantize always outputs fp32. + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP16, DIM); + std::string decoded; + ASSERT_EQ(0, quantizer->dequantize(code.data(), qmeta, &decoded)); + ASSERT_EQ(decoded.size(), DIM * sizeof(float)); + + const float *recon = reinterpret_cast(decoded.data()); + float recon_err = reference_sq_euclidean(orig_fp32.data(), recon, DIM); + float orig_norm = + reference_sq_euclidean(orig_fp32.data(), orig_fp32.data(), DIM); + if (orig_norm > 1e-6f) { + EXPECT_LT(recon_err / orig_norm, 1.0f) + << "recon_err=" << recon_err << " orig_norm=" << orig_norm; + } +} + +// Verify serialize/deserialize round-trip preserves FP16 PQ codes. +TEST(PqInt8Fp16, SerializeDeserialize) { + const size_t DIM = 16; + const size_t NSQ = 4; + const size_t COUNT = 500; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(quantizer); + + auto holder = make_random_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Serialize. + std::string blob; + ASSERT_EQ(0, quantizer->serialize(&blob)); + EXPECT_GT(blob.size(), sizeof(zvec::turbo::QuantizerSerHeader)); + + // Deserialize into a fresh quantizer. + auto q2 = IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_TRUE(q2); + ASSERT_EQ(0, q2->deserialize(blob)); + + // Verify deserialized quantizer reports FP16 input type. + EXPECT_EQ(q2->input_data_type(), DataType::kFp16); + + // Encode the same vector with both and compare codes. + auto iter = holder->create_iterator(); + iter->is_valid(); + std::vector code1(quantizer->quantized_datapoint_vector_length()); + std::vector code2(q2->quantized_datapoint_vector_length()); + quantizer->quantize_data(iter->data(), code1.data()); + q2->quantize_data(iter->data(), code2.data()); + + for (size_t m = 0; m < NSQ; ++m) { + EXPECT_EQ(code1[m], code2[m]) << "m=" << m; + } + + // ADC distances should also match (same codebook → same LUT → same ADC). + size_t lut_len = quantizer->quantized_query_vector_length(); + std::vector lut1(lut_len / sizeof(float)); + std::vector lut2(lut_len / sizeof(float)); + quantizer->quantize_query(iter->data(), lut1.data()); + q2->quantize_query(iter->data(), lut2.data()); + + float adc1 = quantizer->calc_distance_dp_query(code1.data(), lut1.data()); + float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); + EXPECT_NEAR(adc1, adc2, 1e-6f); +} + +// Helper: build a holder with random fp16 vectors with varying norms (Cosine). +static std::shared_ptr> +make_cosine_fp16_holder(size_t count, size_t dim, uint32_t seed = 42) { + auto holder = + std::make_shared>(dim); + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::uniform_real_distribution scale(0.5f, 5.0f); + for (size_t i = 0; i < count; ++i) { + NumericalVector vec(dim); + float s = scale(gen); + for (size_t j = 0; j < dim; ++j) vec[j] = ailego::Float16(dist(gen) * s); + holder->emplace(i + 1, vec); + } + return holder; +} + +// Verify FP16 with Cosine metric: train, encode, ADC distance range. +TEST(PqInt8Fp16, CosineAdcDistance) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + auto quantizer = make_pq_fp16_quantizer(DIM, NSQ, "Cosine"); + ASSERT_TRUE(quantizer); + + // Cosine FP16: extra meta = sizeof(float). + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), + NSQ + sizeof(float)); + + auto holder = make_cosine_fp16_holder(COUNT, DIM); + ASSERT_EQ(0, quantizer->train(holder)); + + // Collect raw vectors (as fp32) and PQ codes. + std::vector> raw_vecs(COUNT); + std::vector> pq_codes(COUNT); + size_t code_len = quantizer->quantized_datapoint_vector_length(); + size_t lut_len = quantizer->quantized_query_vector_length(); + + auto iter = holder->create_iterator(); + for (size_t i = 0; iter->is_valid(); iter->next(), ++i) { + const ailego::Float16 *v = + reinterpret_cast(iter->data()); + raw_vecs[i] = fp16_to_fp32(v, DIM); + pq_codes[i].resize(code_len); + quantizer->quantize_data(iter->data(), pq_codes[i].data()); + } + + // Build LUT for query = first vector (must use FP16 data from holder). + auto iter2 = holder->create_iterator(); + iter2->is_valid(); + std::vector lut(lut_len / sizeof(float)); + quantizer->quantize_query(iter2->data(), lut.data()); + + for (size_t i = 1; i < COUNT; ++i) { + float adc_dist = + quantizer->calc_distance_dp_query(pq_codes[i].data(), lut.data()); + float true_dist = + reference_cosine_distance(raw_vecs[i].data(), raw_vecs[0].data(), DIM); + + // Cosine ADC distance should be in [0, 2] (with some tolerance). + EXPECT_GE(adc_dist, -0.05f) << "i=" << i; + EXPECT_LE(adc_dist, 2.05f) << "i=" << i; + + // PQ approximation: should be roughly correlated. + EXPECT_LT(std::fabs(adc_dist - true_dist), 0.5f) + << "i=" << i << " adc=" << adc_dist << " true=" << true_dist; + } +} + +// Verify FP16 PQ produces ADC distance rankings consistent with FP32 PQ. +// Both quantizers train independent codebooks, so raw codes differ, but the +// relative distance ordering should be largely preserved on the same data. +TEST(PqInt8Fp16, ConsistencyWithFp32) { + const size_t DIM = 32; + const size_t NSQ = 8; + const size_t COUNT = 2000; + + // Train FP32 quantizer. + auto q_fp32 = make_pq_quantizer(DIM, NSQ); + ASSERT_TRUE(q_fp32); + auto holder_fp32 = make_random_holder(COUNT, DIM, 123); + ASSERT_EQ(0, q_fp32->train(holder_fp32)); + + // Build FP16 quantizer with same seed data. + auto q_fp16 = make_pq_fp16_quantizer(DIM, NSQ); + ASSERT_TRUE(q_fp16); + auto holder_fp16 = make_random_fp16_holder(COUNT, DIM, 123); + ASSERT_EQ(0, q_fp16->train(holder_fp16)); + + // Build LUTs for the first vector in each holder. + size_t lut_len = q_fp32->quantized_query_vector_length(); + std::vector lut_fp32(lut_len / sizeof(float)); + std::vector lut_fp16(lut_len / sizeof(float)); + + auto iter32 = holder_fp32->create_iterator(); + auto iter16 = holder_fp16->create_iterator(); + iter32->is_valid(); + iter16->is_valid(); + q_fp32->quantize_query(iter32->data(), lut_fp32.data()); + q_fp16->quantize_query(iter16->data(), lut_fp16.data()); + + // Encode next 200 vectors and collect ADC distances. + size_t code_len = q_fp32->quantized_datapoint_vector_length(); + std::vector code32(code_len), code16(code_len); + + std::vector adc32_vec, adc16_vec; + iter32->next(); + iter16->next(); + for (size_t i = 1; i < 201 && iter32->is_valid() && iter16->is_valid(); + iter32->next(), iter16->next(), ++i) { + q_fp32->quantize_data(iter32->data(), code32.data()); + q_fp16->quantize_data(iter16->data(), code16.data()); + adc32_vec.push_back( + q_fp32->calc_distance_dp_query(code32.data(), lut_fp32.data())); + adc16_vec.push_back( + q_fp16->calc_distance_dp_query(code16.data(), lut_fp16.data())); + } + + // Compute Spearman rank correlation between FP32 and FP16 ADC distances. + size_t n = adc32_vec.size(); + ASSERT_GT(n, 10u); + + // Count concordant vs discordant pairs (Kendall tau). + size_t concordant = 0, discordant = 0; + for (size_t i = 0; i < n; ++i) { + for (size_t j = i + 1; j < n; ++j) { + bool same_order = + (adc32_vec[i] < adc32_vec[j]) == (adc16_vec[i] < adc16_vec[j]); + if (same_order) + ++concordant; + else + ++discordant; + } + } + double tau = static_cast(concordant - discordant) / + static_cast(concordant + discordant); + + // Kendall tau > 0.5 means strong rank correlation (generous threshold; + // FP16 precision loss and different codebooks cause some reordering). + EXPECT_GT(tau, 0.5) << "Kendall tau=" << tau; +}