From 6c6e3e2a4477532ddf2f04fd1300780de804aa02 Mon Sep 17 00:00:00 2001 From: Jianning Wang Date: Fri, 24 Jul 2026 17:24:26 +0800 Subject: [PATCH 1/4] feat: init --- src/turbo/distance/scalar/fp16/cosine.cc | 40 +++ src/turbo/distance/scalar/fp16/cosine.h | 30 ++ .../distance/scalar/fp16/inner_product.cc | 46 +++ .../distance/scalar/fp16/inner_product.h | 32 ++ .../distance/scalar/fp16/squared_euclidean.cc | 45 +++ .../distance/scalar/fp16/squared_euclidean.h | 31 ++ .../scalar/record_quantized_int4/common.h | 58 ++++ .../scalar/record_quantized_int4/cosine.cc | 71 +++++ .../scalar/record_quantized_int4/cosine.h | 32 ++ .../record_quantized_int4/inner_product.cc | 67 +++++ .../record_quantized_int4/inner_product.h | 32 ++ .../squared_euclidean.cc | 63 ++++ .../record_quantized_int4/squared_euclidean.h | 32 ++ .../scalar/record_quantized_int8/common.h | 49 +++ .../scalar/record_quantized_int8/cosine.cc | 70 +++++ .../scalar/record_quantized_int8/cosine.h | 32 ++ .../record_quantized_int8/inner_product.cc | 58 ++++ .../record_quantized_int8/inner_product.h | 32 ++ .../squared_euclidean.cc | 61 ++++ .../record_quantized_int8/squared_euclidean.h | 31 ++ .../fp16_quantizer/fp16_quantizer.cc | 199 +++++++++++++ .../quantizer/fp16_quantizer/fp16_quantizer.h | 126 ++++++++ .../int4_quantizer/int4_quantizer.cc | 225 ++++++++++++++ .../quantizer/int4_quantizer/int4_quantizer.h | 145 +++++++++ .../int8_quantizer/int8_quantizer.cc | 237 +++++++++++++++ .../quantizer/int8_quantizer/int8_quantizer.h | 145 +++++++++ src/turbo/turbo.cc | 103 ++++++- tests/turbo/turbo_fp16_quantizer_test.cc | 258 ++++++++++++++++ tests/turbo/turbo_int4_quantizer_test.cc | 280 ++++++++++++++++++ tests/turbo/turbo_int8_quantizer_test.cc | 269 +++++++++++++++++ 30 files changed, 2896 insertions(+), 3 deletions(-) create mode 100644 src/turbo/distance/scalar/fp16/cosine.cc create mode 100644 src/turbo/distance/scalar/fp16/cosine.h create mode 100644 src/turbo/distance/scalar/fp16/inner_product.cc create mode 100644 src/turbo/distance/scalar/fp16/inner_product.h create mode 100644 src/turbo/distance/scalar/fp16/squared_euclidean.cc create mode 100644 src/turbo/distance/scalar/fp16/squared_euclidean.h create mode 100644 src/turbo/distance/scalar/record_quantized_int4/common.h create mode 100644 src/turbo/distance/scalar/record_quantized_int4/cosine.cc create mode 100644 src/turbo/distance/scalar/record_quantized_int4/cosine.h create mode 100644 src/turbo/distance/scalar/record_quantized_int4/inner_product.cc create mode 100644 src/turbo/distance/scalar/record_quantized_int4/inner_product.h create mode 100644 src/turbo/distance/scalar/record_quantized_int4/squared_euclidean.cc create mode 100644 src/turbo/distance/scalar/record_quantized_int4/squared_euclidean.h create mode 100644 src/turbo/distance/scalar/record_quantized_int8/common.h create mode 100644 src/turbo/distance/scalar/record_quantized_int8/cosine.cc create mode 100644 src/turbo/distance/scalar/record_quantized_int8/cosine.h create mode 100644 src/turbo/distance/scalar/record_quantized_int8/inner_product.cc create mode 100644 src/turbo/distance/scalar/record_quantized_int8/inner_product.h create mode 100644 src/turbo/distance/scalar/record_quantized_int8/squared_euclidean.cc create mode 100644 src/turbo/distance/scalar/record_quantized_int8/squared_euclidean.h create mode 100644 src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc create mode 100644 src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h create mode 100644 src/turbo/quantizer/int4_quantizer/int4_quantizer.cc create mode 100644 src/turbo/quantizer/int4_quantizer/int4_quantizer.h create mode 100644 src/turbo/quantizer/int8_quantizer/int8_quantizer.cc create mode 100644 src/turbo/quantizer/int8_quantizer/int8_quantizer.h create mode 100644 tests/turbo/turbo_fp16_quantizer_test.cc create mode 100644 tests/turbo/turbo_int4_quantizer_test.cc create mode 100644 tests/turbo/turbo_int8_quantizer_test.cc diff --git a/src/turbo/distance/scalar/fp16/cosine.cc b/src/turbo/distance/scalar/fp16/cosine.cc new file mode 100644 index 000000000..ee798cd44 --- /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..3fda42e52 --- /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 (negative inner product after normalization) +// between a single quantized FP16 vector pair. +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..2ec08deb4 --- /dev/null +++ b/src/turbo/distance/scalar/fp16/inner_product.cc @@ -0,0 +1,46 @@ +// 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 +#include + +namespace zvec::turbo::scalar { + +// Compute negated inner product between a single FP16 vector pair. +// Returns -dot(a, b) so that callers can derive cosine distance as 1 + ip. +void inner_product_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + const uint16_t *m = reinterpret_cast(a); + const uint16_t *q = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + sum += zvec::ailego::FloatHelper::ToFP32(m[i]) * + zvec::ailego::FloatHelper::ToFP32(q[i]); + } + + *distance = -sum; +} + +// 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) { + 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..50efcbc6d --- /dev/null +++ b/src/turbo/distance/scalar/fp16/inner_product.h @@ -0,0 +1,32 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute inner product distance between a single quantized FP16 +// vector pair. Returns -dot(a, b) so that callers can derive cosine +// distance as 1 + ip. +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..8cf6402eb --- /dev/null +++ b/src/turbo/distance/scalar/fp16/squared_euclidean.cc @@ -0,0 +1,45 @@ +// 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 +#include +#include + +namespace zvec::turbo::scalar { + +void squared_euclidean_fp16_distance(const void *a, const void *b, size_t dim, + float *distance) { + const uint16_t *m = reinterpret_cast(a); + const uint16_t *q = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + sum += zvec::ailego::MathHelper::SquaredDifference( + zvec::ailego::FloatHelper::ToFP32(m[i]), + zvec::ailego::FloatHelper::ToFP32(q[i])); + } + + *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..971688a0c --- /dev/null +++ b/src/turbo/distance/scalar/fp16/squared_euclidean.h @@ -0,0 +1,31 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute squared euclidean distance between a single quantized 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/record_quantized_int4/common.h b/src/turbo/distance/scalar/record_quantized_int4/common.h new file mode 100644 index 000000000..46187580e --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int4/common.h @@ -0,0 +1,58 @@ +// 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. + +// Shared scalar helpers for record_quantized_int4 distance implementations. +// +// Record-quantized INT4 layout (see core::RecordQuantizer::quantize_record): +// +// [ original_dim / 2 bytes: packed signed int4 elements (lo nibble first) ] +// [ float scale_a ] (a: 1 / scale) +// [ float bias_a ] (b: -bias / scale) +// [ float sum_a ] (s: sum of quantized codes) +// [ float square_sum_a ] (s2: sum of squared codes, euclidean metrics) +// or +// [ int int8_sum ] (sum of raw codes, non-euclidean metrics) +// +// Total tail size: 16 bytes = 32 int4 units. For the Cosine metric an +// additional fp32 norm is appended after the tail (20 bytes = 40 int4 +// units); it is not used for distance computation. +// +// All `dim` arguments below are expressed in int4 units, following the core +// convention (element size in bytes is dim / 2). + +#pragma once + +#include +#include + +namespace zvec::turbo::scalar::internal { + +// Raw integer inner product of two packed signed int4 code arrays holding +// `size` int4 elements (`size` must be even). +inline float ip_int4_scalar(const void *a, const void *b, size_t size) { + const uint8_t *lhs = reinterpret_cast(a); + const uint8_t *rhs = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < (size >> 1); ++i) { + int8_t m_lo = static_cast(lhs[i] << 4) >> 4; + int8_t m_hi = static_cast(lhs[i] & 0xf0) >> 4; + int8_t q_lo = static_cast(rhs[i] << 4) >> 4; + int8_t q_hi = static_cast(rhs[i] & 0xf0) >> 4; + sum += static_cast(m_lo * q_lo + m_hi * q_hi); + } + return sum; +} + +} // namespace zvec::turbo::scalar::internal diff --git a/src/turbo/distance/scalar/record_quantized_int4/cosine.cc b/src/turbo/distance/scalar/record_quantized_int4/cosine.cc new file mode 100644 index 000000000..793628d72 --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int4/cosine.cc @@ -0,0 +1,71 @@ +// 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/record_quantized_int4/cosine.h" +#include +#include "scalar/record_quantized_int4/common.h" + +// Tail layout for record-quantized INT4 cosine vectors: +// +// [ original_dim / 2 bytes: packed signed int4 elements ] +// [ float scale_a ] (ma) +// [ float bias_a ] (mb) +// [ float sum_a ] (ms) +// [ int int8_sum ] (sum of raw codes, non-euclidean layout) +// [ float norm ] (original L2 norm, unused for distance) +// +// The distance returned is the negated dequantized inner product, matching +// core::CosineMinusInnerProduct. Callers normalize it to 1 - cos. + +namespace zvec::turbo::scalar { + +void cosine_int4_distance(const void *a, const void *b, size_t dim, + float *distance) { + // `dim` is the full encoded size in int4 units; the original vector + // occupies dim-40 int4 units. + const int original_dim = static_cast(dim) - 40; + if (original_dim <= 0) { + return; + } + const size_t p = static_cast(original_dim) >> 1; // params offset + float ip = internal::ip_int4_scalar(a, b, original_dim); + + const float *a_tail = + reinterpret_cast(reinterpret_cast(a) + p); + const float *b_tail = + reinterpret_cast(reinterpret_cast(b) + p); + + float ma = a_tail[0]; + float mb = a_tail[1]; + float ms = a_tail[2]; + + float qa = b_tail[0]; + float qb = b_tail[1]; + float qs = b_tail[2]; + + // Dequantize and compute cosine distance: + // cosine_dist = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + // + original_dim * qb * mb) + *distance = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + + static_cast(original_dim) * qb * mb); +} + +void cosine_int4_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances) { + for (size_t i = 0; i < n; ++i) { + cosine_int4_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/record_quantized_int4/cosine.h b/src/turbo/distance/scalar/record_quantized_int4/cosine.h new file mode 100644 index 000000000..0c2854ced --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int4/cosine.h @@ -0,0 +1,32 @@ +// 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 record-quantized INT4 vector +// pair (both sides are expected to be L2-normalized before quantization). +// `dim` is the full encoded size in int4 units (original_dim + 40, where +// the last 4 bytes store the fp32 norm). +void cosine_int4_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of cosine_int4_distance. +void cosine_int4_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/record_quantized_int4/inner_product.cc b/src/turbo/distance/scalar/record_quantized_int4/inner_product.cc new file mode 100644 index 000000000..8e691fe7b --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int4/inner_product.cc @@ -0,0 +1,67 @@ +// 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/record_quantized_int4/inner_product.h" +#include +#include "scalar/record_quantized_int4/common.h" + +namespace zvec::turbo::scalar { + +// Shared implementation over the original int4 dimensionality. +static void minus_inner_product_int4_impl(const void *a, const void *b, + size_t original_dim, + float *distance) { + const size_t p = original_dim >> 1; // params offset in bytes + float ip = internal::ip_int4_scalar(a, b, original_dim); + + const float *a_tail = + reinterpret_cast(reinterpret_cast(a) + p); + const float *b_tail = + reinterpret_cast(reinterpret_cast(b) + p); + + float ma = a_tail[0]; + float mb = a_tail[1]; + float ms = a_tail[2]; + + float qa = b_tail[0]; + float qb = b_tail[1]; + float qs = b_tail[2]; + + // Dequantize and compute the negated inner product: + // ip_dist = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + // + original_dim * qb * mb) + *distance = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + + static_cast(original_dim) * qb * mb); +} + +void inner_product_int4_distance(const void *a, const void *b, size_t dim, + float *distance) { + // `dim` is the full encoded size in int4 units; the original vector + // occupies dim-32 int4 units. + const int original_dim = static_cast(dim) - 32; + if (original_dim <= 0) { + return; + } + minus_inner_product_int4_impl(a, b, original_dim, distance); +} + +void inner_product_int4_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_int4_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/record_quantized_int4/inner_product.h b/src/turbo/distance/scalar/record_quantized_int4/inner_product.h new file mode 100644 index 000000000..bfa9c5362 --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int4/inner_product.h @@ -0,0 +1,32 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute inner product distance (negated dequantized inner product) between +// a single record-quantized INT4 vector pair. `dim` is the full encoded size +// in int4 units (original_dim + 32). +void inner_product_int4_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of inner_product_int4_distance. +void inner_product_int4_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/record_quantized_int4/squared_euclidean.cc b/src/turbo/distance/scalar/record_quantized_int4/squared_euclidean.cc new file mode 100644 index 000000000..2dd2389a0 --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int4/squared_euclidean.cc @@ -0,0 +1,63 @@ +// 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/record_quantized_int4/squared_euclidean.h" +#include +#include "scalar/record_quantized_int4/common.h" + +namespace zvec::turbo::scalar { + +void squared_euclidean_int4_distance(const void *a, const void *b, size_t dim, + float *distance) { + // `dim` is the full encoded size in int4 units; the original vector + // occupies dim-32 int4 units (i.e. (dim-32)/2 bytes). + const int original_dim = static_cast(dim) - 32; + if (original_dim <= 0) { + return; + } + const size_t p = static_cast(original_dim) >> 1; // params offset + float ip = internal::ip_int4_scalar(a, b, original_dim); + + const float *a_tail = + reinterpret_cast(reinterpret_cast(a) + p); + const float *b_tail = + reinterpret_cast(reinterpret_cast(b) + p); + + float ma = a_tail[0]; + float mb = a_tail[1]; + float ms = a_tail[2]; + float ms2 = a_tail[3]; + + float qa = b_tail[0]; + float qb = b_tail[1]; + float qs = b_tail[2]; + float qs2 = b_tail[3]; + + const float sum = qa * qs; + const float sum2 = qa * qa * qs2; + + *distance = ma * ma * ms2 + sum2 - 2 * ma * qa * ip + + (mb - qb) * (mb - qb) * original_dim + + 2 * (mb - qb) * (ms * ma - sum); +} + +void squared_euclidean_int4_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_int4_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/record_quantized_int4/squared_euclidean.h b/src/turbo/distance/scalar/record_quantized_int4/squared_euclidean.h new file mode 100644 index 000000000..4f8e1c2a1 --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int4/squared_euclidean.h @@ -0,0 +1,32 @@ +// 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 record-quantized INT4 +// vector pair. `dim` is the full encoded size in int4 units +// (original_dim + 32). +void squared_euclidean_int4_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of squared_euclidean_int4_distance. +void squared_euclidean_int4_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/record_quantized_int8/common.h b/src/turbo/distance/scalar/record_quantized_int8/common.h new file mode 100644 index 000000000..505f92340 --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int8/common.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. + +// Shared scalar helpers for record_quantized_int8 distance implementations. +// +// Record-quantized INT8 layout (see core::RecordQuantizer::quantize_record): +// +// [ original_dim bytes: int8_t elements ] +// [ float scale_a ] (a: 1 / scale) +// [ float bias_a ] (b: -bias / scale) +// [ float sum_a ] (s: sum of quantized codes) +// [ float square_sum_a ] (s2: sum of squared quantized codes) +// [ int int8_sum ] (sum of raw int8 elements) +// +// Total tail size: 4 floats + 1 int = 20 bytes. For the Cosine metric an +// additional fp32 norm is appended after the tail (total 24 bytes); it is +// not used for distance computation. + +#pragma once + +#include +#include + +namespace zvec::turbo::scalar::internal { + +// Raw integer inner product of two int8 code arrays of length `size`. +inline float ip_int8_scalar(const void *a, const void *b, size_t size) { + const int8_t *lhs = reinterpret_cast(a); + const int8_t *rhs = reinterpret_cast(b); + + float sum = 0.0f; + for (size_t i = 0; i < size; ++i) { + sum += static_cast(lhs[i] * rhs[i]); + } + return sum; +} + +} // namespace zvec::turbo::scalar::internal diff --git a/src/turbo/distance/scalar/record_quantized_int8/cosine.cc b/src/turbo/distance/scalar/record_quantized_int8/cosine.cc new file mode 100644 index 000000000..1390ee7ec --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int8/cosine.cc @@ -0,0 +1,70 @@ +// 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/record_quantized_int8/cosine.h" +#include +#include "scalar/record_quantized_int8/common.h" + +// Tail layout for record-quantized INT8 cosine vectors: +// +// [ original_dim bytes: int8_t elements ] +// [ float scale_a ] (ma) +// [ float bias_a ] (mb) +// [ float sum_a ] (ms) +// [ float square_sum_a ] (ms2) +// [ int int8_sum ] +// [ float norm ] (original L2 norm, unused for distance) +// +// The distance returned is the negated dequantized inner product, matching +// avx512_vnni::cosine_int8_distance. Callers normalize it to 1 - cos. + +namespace zvec::turbo::scalar { + +void cosine_int8_distance(const void *a, const void *b, size_t dim, + float *distance) { + // `dim` is the full encoded size; the original vector occupies dim-24 bytes. + const int original_dim = static_cast(dim) - 24; + if (original_dim <= 0) { + return; + } + float ip = internal::ip_int8_scalar(a, b, original_dim); + + const float *a_tail = reinterpret_cast( + reinterpret_cast(a) + original_dim); + const float *b_tail = reinterpret_cast( + reinterpret_cast(b) + original_dim); + + float ma = a_tail[0]; + float mb = a_tail[1]; + float ms = a_tail[2]; + + float qa = b_tail[0]; + float qb = b_tail[1]; + float qs = b_tail[2]; + + // Dequantize and compute cosine distance: + // cosine_dist = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + // + original_dim * qb * mb) + *distance = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + + static_cast(original_dim) * qb * mb); +} + +void cosine_int8_batch_distance(const void *const *vectors, const void *query, + size_t n, size_t dim, float *distances) { + for (size_t i = 0; i < n; ++i) { + cosine_int8_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/record_quantized_int8/cosine.h b/src/turbo/distance/scalar/record_quantized_int8/cosine.h new file mode 100644 index 000000000..fc7fe5f6c --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int8/cosine.h @@ -0,0 +1,32 @@ +// 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 record-quantized INT8 vector +// pair (both sides are expected to be L2-normalized before quantization). +// `dim` is the full encoded size (original_dim + 24, where the last 4 bytes +// store the fp32 norm). +void cosine_int8_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of cosine_int8_distance. +void cosine_int8_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/record_quantized_int8/inner_product.cc b/src/turbo/distance/scalar/record_quantized_int8/inner_product.cc new file mode 100644 index 000000000..76fd46302 --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int8/inner_product.cc @@ -0,0 +1,58 @@ +// 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/record_quantized_int8/inner_product.h" +#include +#include "scalar/record_quantized_int8/common.h" + +namespace zvec::turbo::scalar { + +void inner_product_int8_distance(const void *a, const void *b, size_t dim, + float *distance) { + // `dim` is the full encoded size; the original vector occupies dim-20 bytes. + const int original_dim = static_cast(dim) - 20; + if (original_dim <= 0) { + return; + } + float ip = internal::ip_int8_scalar(a, b, original_dim); + + const float *a_tail = reinterpret_cast( + reinterpret_cast(a) + original_dim); + const float *b_tail = reinterpret_cast( + reinterpret_cast(b) + original_dim); + + float ma = a_tail[0]; + float mb = a_tail[1]; + float ms = a_tail[2]; + + float qa = b_tail[0]; + float qb = b_tail[1]; + float qs = b_tail[2]; + + // Dequantize and compute the negated inner product: + // ip_dist = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + // + original_dim * qb * mb) + *distance = -(ma * qa * ip + mb * qa * qs + qb * ma * ms + + static_cast(original_dim) * qb * mb); +} + +void inner_product_int8_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_int8_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/record_quantized_int8/inner_product.h b/src/turbo/distance/scalar/record_quantized_int8/inner_product.h new file mode 100644 index 000000000..f73460de2 --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int8/inner_product.h @@ -0,0 +1,32 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute inner product distance (negated dequantized inner product) between +// a single record-quantized INT8 vector pair. `dim` is the full encoded size +// (original_dim + 20). +void inner_product_int8_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of inner_product_int8_distance. +void inner_product_int8_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/record_quantized_int8/squared_euclidean.cc b/src/turbo/distance/scalar/record_quantized_int8/squared_euclidean.cc new file mode 100644 index 000000000..f9ed83bce --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int8/squared_euclidean.cc @@ -0,0 +1,61 @@ +// 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/record_quantized_int8/squared_euclidean.h" +#include +#include "scalar/record_quantized_int8/common.h" + +namespace zvec::turbo::scalar { + +void squared_euclidean_int8_distance(const void *a, const void *b, size_t dim, + float *distance) { + // `dim` is the full encoded size; the original vector occupies dim-20 bytes. + const int original_dim = static_cast(dim) - 20; + if (original_dim <= 0) { + return; + } + float ip = internal::ip_int8_scalar(a, b, original_dim); + + const float *a_tail = reinterpret_cast( + reinterpret_cast(a) + original_dim); + const float *b_tail = reinterpret_cast( + reinterpret_cast(b) + original_dim); + + float ma = a_tail[0]; + float mb = a_tail[1]; + float ms = a_tail[2]; + float ms2 = a_tail[3]; + + float qa = b_tail[0]; + float qb = b_tail[1]; + float qs = b_tail[2]; + float qs2 = b_tail[3]; + + const float sum = qa * qs; + const float sum2 = qa * qa * qs2; + + *distance = ma * ma * ms2 + sum2 - 2 * ma * qa * ip + + (mb - qb) * (mb - qb) * original_dim + + 2 * (mb - qb) * (ms * ma - sum); +} + +void squared_euclidean_int8_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_int8_distance(vectors[i], query, dim, &distances[i]); + } +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/record_quantized_int8/squared_euclidean.h b/src/turbo/distance/scalar/record_quantized_int8/squared_euclidean.h new file mode 100644 index 000000000..7cc965a7a --- /dev/null +++ b/src/turbo/distance/scalar/record_quantized_int8/squared_euclidean.h @@ -0,0 +1,31 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include + +namespace zvec::turbo::scalar { + +// Compute squared euclidean distance between a single record-quantized INT8 +// vector pair. `dim` is the full encoded size (original_dim + 20). +void squared_euclidean_int8_distance(const void *a, const void *b, size_t dim, + float *distance); + +// Batch version of squared_euclidean_int8_distance. +void squared_euclidean_int8_batch_distance(const void *const *vectors, + const void *query, size_t n, + size_t dim, float *distances); + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc new file mode 100644 index 000000000..7ff6517ef --- /dev/null +++ b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc @@ -0,0 +1,199 @@ +// 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/fp16_quantizer/fp16_quantizer.h" +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace turbo { + +int Fp16Quantizer::init(const IndexMeta &meta, + const ailego::Params & /*params*/) { + meta_ = meta; + + meta_.set_meta(IndexMeta::DataType::DT_FP16, meta.dimension()); + + original_dim_ = meta.dimension(); + auto metric_name = meta.metric_name(); + if (metric_name == "Cosine") { + extra_meta_size_ = EXTRA_META_SIZE_COSINE; + meta_.set_extra_meta_size(extra_meta_size_); + } + + // Cache the distance dispatch for the new Quantizer interface. + dp_query_func_ = + get_distance_func(metric_from_name(metric_name), DataType::kFp16, + QuantizeType::kFp16, CpuArchType::kAuto); + dp_query_batch_func_ = + get_batch_distance_func(metric_from_name(metric_name), DataType::kFp16, + QuantizeType::kFp16, CpuArchType::kAuto); + + return 0; +} + +int Fp16Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (qmeta.unit_size() != sizeof(float)) { + return kErrUnsupported; + } + + // qmeta.dimension() may be the inflated (data + extras) dimension when the + // caller uses meta_.dimension() directly. Use the raw original dim we + // recorded at init() to avoid over-reading the query. + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + size_t byte_size = raw_dim * sizeof(uint16_t) + extra_meta_size_; + out->resize(byte_size); + uint16_t *out_buf = reinterpret_cast(&(*out)[0]); + + if (meta_.metric_name() == "Cosine") { + // L2-normalize the vector before converting to fp16 and store the norm + // at the end so the original vector can be reconstructed during + // dequantize. + std::vector buf(raw_dim); + std::memcpy(buf.data(), query, raw_dim * sizeof(float)); + float norm = 0.0f; + ailego::Normalizer::L2(buf.data(), raw_dim, &norm); + ailego::FloatHelper::ToFP16(buf.data(), raw_dim, out_buf); + std::memcpy( + reinterpret_cast(&(*out)[0]) + raw_dim * sizeof(uint16_t), + &norm, extra_meta_size_); + } else { + ailego::FloatHelper::ToFP16(reinterpret_cast(query), raw_dim, + out_buf); + } + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_FP16, raw_dim, + static_cast(type_), extra_meta_size_); + + return 0; +} + +int Fp16Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + size_t byte_size = raw_dim * sizeof(float); + + out->resize(byte_size); + const uint16_t *in_buf = reinterpret_cast(in); + float *out_buf = reinterpret_cast(&(*out)[0]); + ailego::FloatHelper::ToFP32(in_buf, raw_dim, out_buf); + + if (meta_.metric_name() == "Cosine") { + // Denormalize the vector using the stored norm. + float norm = 0.0f; + std::memcpy( + &norm, + reinterpret_cast(in) + raw_dim * sizeof(uint16_t), + extra_meta_size_); + for (size_t i = 0; i < raw_dim; ++i) { + out_buf[i] *= norm; + } + } + return 0; +} + +DistanceImpl Fp16Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + auto metric = metric_from_name(meta_.metric_name()); + auto func = get_distance_func(metric, DataType::kFp16, QuantizeType::kFp16, + CpuArchType::kAuto); + if (!func) { + return DistanceImpl{}; + } + auto batch_func = get_batch_distance_func( + metric, DataType::kFp16, QuantizeType::kFp16, CpuArchType::kAuto); + + // The query is assumed to be already quantized — copy it directly. + std::string quantized_query(static_cast(query), + qmeta.element_size()); + return DistanceImpl(std::move(func), std::move(batch_func), + std::move(quantized_query), original_dim_); +} + +void Fp16Quantizer::quantize_one(const void *input, void *output) const { + uint16_t *out_buf = reinterpret_cast(output); + + if (meta_.metric_name() == "Cosine") { + // L2-normalize before converting and store the norm at the end. + std::vector buf(original_dim_); + std::memcpy(buf.data(), input, original_dim_ * sizeof(float)); + float norm = 0.0f; + ailego::Normalizer::L2(buf.data(), original_dim_, &norm); + ailego::FloatHelper::ToFP16(buf.data(), original_dim_, out_buf); + std::memcpy(reinterpret_cast(output) + + static_cast(original_dim_) * sizeof(uint16_t), + &norm, extra_meta_size_); + } else { + ailego::FloatHelper::ToFP16(reinterpret_cast(input), + original_dim_, out_buf); + } +} + +float Fp16Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + float d = 0.0f; + if (dp_query_func_) { + dp_query_func_(dp, query, original_dim_, &d); + } + return d; +} + +void Fp16Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const { + if (dp_query_batch_func_) { + dp_query_batch_func_(const_cast(dp_list), query, + static_cast(dp_num), original_dim_, dist_list); + return; + } + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = calc_distance_dp_query(dp_list[i], query); + } +} + +float Fp16Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + return calc_distance_dp_query(dp, buf.data()); +} + +void Fp16Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + calc_distance_dp_query_batch(dp_list, dp_num, buf.data(), dist_list); +} + +float Fp16Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + return calc_distance_dp_query(dp1, dp2); +} + +INDEX_FACTORY_REGISTER_QUANTIZER(Fp16Quantizer); + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h new file mode 100644 index 000000000..6b87b86cd --- /dev/null +++ b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h @@ -0,0 +1,126 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +class Fp16Quantizer : public Quantizer { + public: + Fp16Quantizer() { + type_ = QuantizeType::kFp16; + } + + virtual ~Fp16Quantizer() {} + + public: + int init(const core::IndexMeta &meta, const ailego::Params ¶ms) override; + + const core::IndexMeta &meta(void) const override { + return meta_; + } + + DataType input_data_type() const override { + return DataType::kFp32; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return false; + } + + int train(core::IndexHolder::Pointer /*holder*/) override { + return 0; + } + + size_t quantized_datapoint_vector_length() const override { + return quantized_length(); + } + + size_t quantized_query_vector_length() const override { + return quantized_length(); + } + + void quantize_data(const void *input, void *output) const override { + quantize_one(input, output); + } + + void quantize_query(const void *input, void *output) const override { + quantize_one(input, output); + } + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const core::IndexQueryMeta &qmeta, + std::string *out, core::IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const core::IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const core::IndexQueryMeta &qmeta) const override; + + private: + //! Byte length of a quantized vector (fp16 data + extra meta). + size_t quantized_length() const { + return static_cast(original_dim_) * sizeof(uint16_t) + + extra_meta_size_; + } + + //! Quantize a single fp32 vector into a caller-provided buffer of + //! quantized_length() bytes. + void quantize_one(const void *input, void *output) const; + + static constexpr uint32_t EXTRA_META_SIZE_COSINE = 4; + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + + //! Cached distance dispatch (bound in init()). + DistanceFunc dp_query_func_{}; + BatchDistanceFunc dp_query_batch_func_{}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/int4_quantizer/int4_quantizer.cc b/src/turbo/quantizer/int4_quantizer/int4_quantizer.cc new file mode 100644 index 000000000..5bb76d47a --- /dev/null +++ b/src/turbo/quantizer/int4_quantizer/int4_quantizer.cc @@ -0,0 +1,225 @@ +// 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/int4_quantizer/int4_quantizer.h" +#include +#include +#include +#include +#include +#include +#include "core/quantizer/record_quantizer.h" + +namespace zvec { +namespace turbo { + +int Int4Quantizer::init(const IndexMeta &meta, + const ailego::Params & /*params*/) { + if (meta.dimension() % 2 != 0) { + LOG_ERROR("Unsupported dimension %u for INT4 type", meta.dimension()); + return kErrUnsupported; + } + + meta_ = meta; + + meta_.set_meta(IndexMeta::DataType::DT_INT4, meta.dimension()); + + original_dim_ = meta.dimension(); + auto metric_name = meta.metric_name(); + is_cosine_ = (metric_name == "Cosine"); + is_euclidean_ = + (metric_name == "SquaredEuclidean" || metric_name == "Euclidean" || + metric_name == "MipsSquaredEuclidean"); + + extra_meta_size_ = RECORD_TAIL_SIZE; + if (is_cosine_) { + // The raw kernels return -cos; offset to 1 - cos for a non-negative, + // monotonically equivalent distance (same convention as Fp32Quantizer). + distance_offset_ = 1.0f; + extra_meta_size_ += EXTRA_META_SIZE_COSINE; + } + meta_.set_extra_meta_size(extra_meta_size_); + + // Distance kernels take the full encoded size in int4 units (each extra + // byte accounts for two int4 units). + dist_dim_ = static_cast(original_dim_) + 2 * extra_meta_size_; + + // Cache the distance dispatch for the new Quantizer interface. + auto metric = metric_from_name(metric_name); + dp_query_func_ = get_distance_func(metric, DataType::kInt4, + QuantizeType::kRecord, CpuArchType::kAuto); + dp_query_batch_func_ = get_batch_distance_func( + metric, DataType::kInt4, QuantizeType::kRecord, CpuArchType::kAuto); + + return 0; +} + +void Int4Quantizer::quantize_one(const void *input, void *output) const { + const float *vec = reinterpret_cast(input); + + std::vector normalized; + float norm = 0.0f; + if (is_cosine_) { + // L2-normalize before quantization; the norm is stored after the record + // tail so the original vector can be reconstructed during dequantize. + normalized.assign(vec, vec + original_dim_); + ailego::Normalizer::L2(normalized.data(), original_dim_, &norm); + vec = normalized.data(); + } + + RecordQuantizer::quantize_record( + vec, original_dim_, IndexMeta::DataType::DT_INT4, is_euclidean_, output); + + if (is_cosine_) { + std::memcpy(reinterpret_cast(output) + original_dim_ / 2 + + RECORD_TAIL_SIZE, + &norm, EXTRA_META_SIZE_COSINE); + } +} + +int Int4Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (qmeta.unit_size() != sizeof(float)) { + return kErrUnsupported; + } + + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + if (raw_dim != original_dim_) { + return kErrUnsupported; + } + + out->resize(quantized_length()); + quantize_one(query, &(*out)[0]); + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_INT4, raw_dim, + static_cast(type_), extra_meta_size_); + + return 0; +} + +int Int4Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + + out->resize(raw_dim * sizeof(float)); + float *out_buf = reinterpret_cast(&(*out)[0]); + RecordQuantizer::unquantize_record(in, raw_dim, IndexMeta::DataType::DT_INT4, + out_buf); + + if (is_cosine_) { + // Denormalize the vector using the stored norm. + float norm = 0.0f; + std::memcpy( + &norm, + reinterpret_cast(in) + raw_dim / 2 + RECORD_TAIL_SIZE, + EXTRA_META_SIZE_COSINE); + for (size_t i = 0; i < raw_dim; ++i) { + out_buf[i] *= norm; + } + } + return 0; +} + +DistanceImpl Int4Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + DistanceFunc func = dp_query_func_; + if (!func) { + return DistanceImpl{}; + } + BatchDistanceFunc batch_func = dp_query_batch_func_; + + if (distance_offset_ != 0.0f) { + float offset = distance_offset_; + DistanceFunc base = std::move(func); + func = [base, offset](const void *a, const void *b, size_t dim, + float *out) { + base(a, b, dim, out); + *out += offset; + }; + if (batch_func) { + BatchDistanceFunc batch_base = std::move(batch_func); + batch_func = [batch_base, offset](const void **m, const void *q, + size_t num, size_t dim, float *out) { + batch_base(m, q, num, dim, out); + for (size_t i = 0; i < num; ++i) { + out[i] += offset; + } + }; + } + } + + // The query is assumed to be already quantized — copy it directly. + std::string quantized_query(static_cast(query), + qmeta.element_size()); + return DistanceImpl(std::move(func), std::move(batch_func), + std::move(quantized_query), dist_dim_); +} + +float Int4Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + float d = 0.0f; + if (dp_query_func_) { + dp_query_func_(dp, query, dist_dim_, &d); + d += distance_offset_; + } + return d; +} + +void Int4Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const { + if (dp_query_batch_func_) { + dp_query_batch_func_(const_cast(dp_list), query, + static_cast(dp_num), dist_dim_, dist_list); + if (distance_offset_ != 0.0f) { + for (int i = 0; i < dp_num; ++i) { + dist_list[i] += distance_offset_; + } + } + return; + } + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = calc_distance_dp_query(dp_list[i], query); + } +} + +float Int4Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + return calc_distance_dp_query(dp, buf.data()); +} + +void Int4Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + calc_distance_dp_query_batch(dp_list, dp_num, buf.data(), dist_list); +} + +float Int4Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + return calc_distance_dp_query(dp1, dp2); +} + +INDEX_FACTORY_REGISTER_QUANTIZER(Int4Quantizer); + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/int4_quantizer/int4_quantizer.h b/src/turbo/quantizer/int4_quantizer/int4_quantizer.h new file mode 100644 index 000000000..435e8ebcd --- /dev/null +++ b/src/turbo/quantizer/int4_quantizer/int4_quantizer.h @@ -0,0 +1,145 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +/*! Record-quantized INT4 quantizer. + * + * Each vector is quantized independently with a per-record affine transform + * (see core::RecordQuantizer::quantize_record). Layout: + * + * [ dim / 2 bytes packed int4 codes ] + * [ 16-byte tail (a, b, s, s2 | int8_sum) ] + * [ 4-byte fp32 norm ] (Cosine only) + * + * The dimension must be even. + */ +class Int4Quantizer : public Quantizer { + public: + Int4Quantizer() { + type_ = QuantizeType::kRecord; + } + + virtual ~Int4Quantizer() {} + + public: + int init(const core::IndexMeta &meta, const ailego::Params ¶ms) override; + + const core::IndexMeta &meta(void) const override { + return meta_; + } + + DataType input_data_type() const override { + return DataType::kFp32; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return false; + } + + int train(core::IndexHolder::Pointer /*holder*/) override { + return 0; + } + + size_t quantized_datapoint_vector_length() const override { + return quantized_length(); + } + + size_t quantized_query_vector_length() const override { + return quantized_length(); + } + + void quantize_data(const void *input, void *output) const override { + quantize_one(input, output); + } + + void quantize_query(const void *input, void *output) const override { + quantize_one(input, output); + } + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const core::IndexQueryMeta &qmeta, + std::string *out, core::IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const core::IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const core::IndexQueryMeta &qmeta) const override; + + private: + //! Byte length of a quantized vector (packed int4 codes + extra meta). + size_t quantized_length() const { + return static_cast(original_dim_) / 2 + extra_meta_size_; + } + + //! Quantize a single fp32 vector into a caller-provided buffer of + //! quantized_length() bytes. + void quantize_one(const void *input, void *output) const; + + //! Record tail: 4 floats (a, b, s, s2 | int8_sum). + static constexpr uint32_t RECORD_TAIL_SIZE = 4 * sizeof(float); + //! Extra fp32 norm appended for the Cosine metric. + static constexpr uint32_t EXTRA_META_SIZE_COSINE = 4; + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + bool is_cosine_{false}; + bool is_euclidean_{false}; + //! Full encoded size in int4 units handed to the distance kernels. + size_t dist_dim_{0}; + //! Offset added to raw kernel outputs (1.0 for Cosine: -cos -> 1 - cos). + float distance_offset_{0.0f}; + + //! Cached distance dispatch (bound in init()). + DistanceFunc dp_query_func_{}; + BatchDistanceFunc dp_query_batch_func_{}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/int8_quantizer/int8_quantizer.cc b/src/turbo/quantizer/int8_quantizer/int8_quantizer.cc new file mode 100644 index 000000000..b6c8bfeb4 --- /dev/null +++ b/src/turbo/quantizer/int8_quantizer/int8_quantizer.cc @@ -0,0 +1,237 @@ +// 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/int8_quantizer/int8_quantizer.h" +#include +#include +#include +#include +#include +#include +#include "core/quantizer/record_quantizer.h" + +namespace zvec { +namespace turbo { + +int Int8Quantizer::init(const IndexMeta &meta, + const ailego::Params & /*params*/) { + meta_ = meta; + + meta_.set_meta(IndexMeta::DataType::DT_INT8, meta.dimension()); + + original_dim_ = meta.dimension(); + auto metric_name = meta.metric_name(); + is_cosine_ = (metric_name == "Cosine"); + is_euclidean_ = + (metric_name == "SquaredEuclidean" || metric_name == "Euclidean" || + metric_name == "MipsSquaredEuclidean"); + + extra_meta_size_ = RECORD_TAIL_SIZE; + if (is_cosine_) { + // The raw kernels return -cos; offset to 1 - cos for a non-negative, + // monotonically equivalent distance (same convention as Fp32Quantizer). + distance_offset_ = 1.0f; + extra_meta_size_ += EXTRA_META_SIZE_COSINE; + } + meta_.set_extra_meta_size(extra_meta_size_); + + // Distance kernels take the full encoded size in bytes. + dist_dim_ = static_cast(original_dim_) + extra_meta_size_; + + // Cache the distance dispatch for the new Quantizer interface. + auto metric = metric_from_name(metric_name); + dp_query_func_ = get_distance_func(metric, DataType::kInt8, + QuantizeType::kRecord, CpuArchType::kAuto); + dp_query_batch_func_ = get_batch_distance_func( + metric, DataType::kInt8, QuantizeType::kRecord, CpuArchType::kAuto); + dp_query_preprocess_func_ = get_query_preprocess_func( + metric, DataType::kInt8, QuantizeType::kRecord, CpuArchType::kAuto); + + return 0; +} + +void Int8Quantizer::quantize_one(const void *input, void *output) const { + const float *vec = reinterpret_cast(input); + + std::vector normalized; + float norm = 0.0f; + if (is_cosine_) { + // L2-normalize before quantization; the norm is stored after the record + // tail so the original vector can be reconstructed during dequantize. + normalized.assign(vec, vec + original_dim_); + ailego::Normalizer::L2(normalized.data(), original_dim_, &norm); + vec = normalized.data(); + } + + RecordQuantizer::quantize_record( + vec, original_dim_, IndexMeta::DataType::DT_INT8, is_euclidean_, output); + + if (is_cosine_) { + std::memcpy( + reinterpret_cast(output) + original_dim_ + RECORD_TAIL_SIZE, + &norm, EXTRA_META_SIZE_COSINE); + } +} + +int Int8Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, + std::string *out, IndexQueryMeta *ometa) const { + if (qmeta.unit_size() != sizeof(float)) { + return kErrUnsupported; + } + + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + if (raw_dim != original_dim_) { + return kErrUnsupported; + } + + out->resize(quantized_length()); + quantize_one(query, &(*out)[0]); + + *ometa = qmeta; + ometa->set_meta(IndexMeta::DataType::DT_INT8, raw_dim, + static_cast(type_), extra_meta_size_); + + return 0; +} + +int Int8Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, + std::string *out) const { + size_t raw_dim = (original_dim_ != 0 && qmeta.dimension() >= original_dim_) + ? original_dim_ + : qmeta.dimension(); + + out->resize(raw_dim * sizeof(float)); + float *out_buf = reinterpret_cast(&(*out)[0]); + RecordQuantizer::unquantize_record(in, raw_dim, IndexMeta::DataType::DT_INT8, + out_buf); + + if (is_cosine_) { + // Denormalize the vector using the stored norm. + float norm = 0.0f; + std::memcpy( + &norm, + reinterpret_cast(in) + raw_dim + RECORD_TAIL_SIZE, + EXTRA_META_SIZE_COSINE); + for (size_t i = 0; i < raw_dim; ++i) { + out_buf[i] *= norm; + } + } + return 0; +} + +DistanceImpl Int8Quantizer::distance(const void *query, + const IndexQueryMeta &qmeta) const { + DistanceFunc func = dp_query_func_; + if (!func) { + return DistanceImpl{}; + } + + // Batch kernels that require a preprocessed (uint8-shifted) query cannot + // share the stored query with the single-pair kernel; fall back to the + // scalar path inside DistanceImpl in that case. + BatchDistanceFunc batch_func; + if (!dp_query_preprocess_func_) { + batch_func = dp_query_batch_func_; + } + + if (distance_offset_ != 0.0f) { + float offset = distance_offset_; + DistanceFunc base = std::move(func); + func = [base, offset](const void *a, const void *b, size_t dim, + float *out) { + base(a, b, dim, out); + *out += offset; + }; + if (batch_func) { + BatchDistanceFunc batch_base = std::move(batch_func); + batch_func = [batch_base, offset](const void **m, const void *q, + size_t num, size_t dim, float *out) { + batch_base(m, q, num, dim, out); + for (size_t i = 0; i < num; ++i) { + out[i] += offset; + } + }; + } + } + + // The query is assumed to be already quantized — copy it directly. + std::string quantized_query(static_cast(query), + qmeta.element_size()); + return DistanceImpl(std::move(func), std::move(batch_func), + std::move(quantized_query), dist_dim_); +} + +float Int8Quantizer::calc_distance_dp_query(const void *dp, + const void *query) const { + float d = 0.0f; + if (dp_query_func_) { + dp_query_func_(dp, query, dist_dim_, &d); + d += distance_offset_; + } + return d; +} + +void Int8Quantizer::calc_distance_dp_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const { + if (dp_query_batch_func_) { + if (dp_query_preprocess_func_) { + // The batch kernel expects the query shifted to uint8; preprocess a + // private copy to keep the caller's buffer intact. + std::string buf(static_cast(query), quantized_length()); + dp_query_preprocess_func_(&buf[0], dist_dim_); + dp_query_batch_func_(const_cast(dp_list), buf.data(), + static_cast(dp_num), dist_dim_, dist_list); + } else { + dp_query_batch_func_(const_cast(dp_list), query, + static_cast(dp_num), dist_dim_, dist_list); + } + if (distance_offset_ != 0.0f) { + for (int i = 0; i < dp_num; ++i) { + dist_list[i] += distance_offset_; + } + } + return; + } + for (int i = 0; i < dp_num; ++i) { + dist_list[i] = calc_distance_dp_query(dp_list[i], query); + } +} + +float Int8Quantizer::calc_distance_dp_query_unquantized( + const void *dp, const void *query) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + return calc_distance_dp_query(dp, buf.data()); +} + +void Int8Quantizer::calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const { + std::string buf(quantized_length(), '\0'); + quantize_one(query, &buf[0]); + calc_distance_dp_query_batch(dp_list, dp_num, buf.data(), dist_list); +} + +float Int8Quantizer::calc_distance_dp_dp(const void *dp1, + const void *dp2) const { + return calc_distance_dp_query(dp1, dp2); +} + +INDEX_FACTORY_REGISTER_QUANTIZER(Int8Quantizer); + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/quantizer/int8_quantizer/int8_quantizer.h b/src/turbo/quantizer/int8_quantizer/int8_quantizer.h new file mode 100644 index 000000000..d8003c5d2 --- /dev/null +++ b/src/turbo/quantizer/int8_quantizer/int8_quantizer.h @@ -0,0 +1,145 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include +#include "quantizer/quantizer.h" + +namespace zvec { +namespace turbo { + +using namespace zvec::core; + +/*! Record-quantized INT8 quantizer. + * + * Each vector is quantized independently with a per-record affine transform + * (see core::RecordQuantizer::quantize_record). Layout: + * + * [ dim bytes int8 codes ][ 20-byte tail (a, b, s, s2, int8_sum) ] + * [ 4-byte fp32 norm ] (Cosine only) + */ +class Int8Quantizer : public Quantizer { + public: + Int8Quantizer() { + type_ = QuantizeType::kRecord; + } + + virtual ~Int8Quantizer() {} + + public: + int init(const core::IndexMeta &meta, const ailego::Params ¶ms) override; + + const core::IndexMeta &meta(void) const override { + return meta_; + } + + DataType input_data_type() const override { + return DataType::kFp32; + } + + QuantizeType type() const override { + return type_; + } + + int dim() const override { + return static_cast(original_dim_); + } + + bool require_train() const override { + return false; + } + + int train(core::IndexHolder::Pointer /*holder*/) override { + return 0; + } + + size_t quantized_datapoint_vector_length() const override { + return quantized_length(); + } + + size_t quantized_query_vector_length() const override { + return quantized_length(); + } + + void quantize_data(const void *input, void *output) const override { + quantize_one(input, output); + } + + void quantize_query(const void *input, void *output) const override { + quantize_one(input, output); + } + + float calc_distance_dp_query(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_dp_query_unquantized(const void *dp, + const void *query) const override; + + void calc_distance_dp_query_batch_unquantized( + const void *const *dp_list, int dp_num, const void *query, + float *dist_list) const override; + + float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + + int quantize(const void *query, const core::IndexQueryMeta &qmeta, + std::string *out, core::IndexQueryMeta *ometa) const override; + + int dequantize(const void *in, const core::IndexQueryMeta &qmeta, + std::string *out) const override; + + DistanceImpl distance(const void *query, + const core::IndexQueryMeta &qmeta) const override; + + private: + //! Byte length of a quantized vector (int8 codes + extra meta). + size_t quantized_length() const { + return static_cast(original_dim_) + extra_meta_size_; + } + + //! Quantize a single fp32 vector into a caller-provided buffer of + //! quantized_length() bytes. + void quantize_one(const void *input, void *output) const; + + //! Record tail: 4 floats (a, b, s, s2) + 1 int32 (int8_sum). + static constexpr uint32_t RECORD_TAIL_SIZE = 4 * sizeof(float) + sizeof(int); + //! Extra fp32 norm appended for the Cosine metric. + static constexpr uint32_t EXTRA_META_SIZE_COSINE = 4; + + IndexMeta meta_{}; + uint32_t original_dim_{0}; + bool is_cosine_{false}; + bool is_euclidean_{false}; + //! Full encoded size in bytes handed to the distance kernels. + size_t dist_dim_{0}; + //! Offset added to raw kernel outputs (1.0 for Cosine: -cos -> 1 - cos). + float distance_offset_{0.0f}; + + //! Cached distance dispatch (bound in init()). + DistanceFunc dp_query_func_{}; + BatchDistanceFunc dp_query_batch_func_{}; + //! Non-null when the batch kernel requires a preprocessed (uint8-shifted) + //! query (AVX512-VNNI path). + QueryPreprocessFunc dp_query_preprocess_func_{nullptr}; +}; + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 5a6f33d49..1b4330890 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -22,9 +22,18 @@ #include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/squared_euclidean.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/record_quantized_int4/cosine.h" +#include "scalar/record_quantized_int4/inner_product.h" +#include "scalar/record_quantized_int4/squared_euclidean.h" +#include "scalar/record_quantized_int8/cosine.h" +#include "scalar/record_quantized_int8/inner_product.h" +#include "scalar/record_quantized_int8/squared_euclidean.h" #include "scalar/rotate/fht/fht.h" #include "sse/rotate/fht/fht.h" @@ -53,8 +62,38 @@ 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::kCosine) { + return scalar::cosine_fp16_distance; + } + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp16_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp16_distance; + } + } + return nullptr; + } + if (data_type == DataType::kInt4) { + if (quantize_type == QuantizeType::kRecord) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_int4_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_int4_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_int4_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { - if (quantize_type == QuantizeType::kDefault) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kRecord) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { @@ -64,6 +103,19 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, return avx512_vnni::cosine_int8_distance; } } + // Scalar fallbacks are only exposed for the kRecord quantize type to + // keep the historical kDefault dispatch behavior unchanged. + if (quantize_type == QuantizeType::kRecord) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_int8_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_int8_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_int8_distance; + } + } } if (quantize_type == QuantizeType::kUniform) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && @@ -96,8 +148,38 @@ 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::kCosine) { + return scalar::cosine_fp16_batch_distance; + } + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_fp16_batch_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_fp16_batch_distance; + } + } + return nullptr; + } + if (data_type == DataType::kInt4) { + if (quantize_type == QuantizeType::kRecord) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_int4_batch_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_int4_batch_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_int4_batch_distance; + } + } + return nullptr; + } if (data_type == DataType::kInt8) { - if (quantize_type == QuantizeType::kDefault) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kRecord) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { @@ -107,6 +189,20 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type, return avx512_vnni::cosine_int8_batch_distance; } } + // Scalar fallbacks are only exposed for the kRecord quantize type to + // keep the historical kDefault dispatch behavior unchanged. The scalar + // batch kernels take a plain int8 query (no preprocessing required). + if (quantize_type == QuantizeType::kRecord) { + if (metric_type == MetricType::kSquaredEuclidean) { + return scalar::squared_euclidean_int8_batch_distance; + } + if (metric_type == MetricType::kCosine) { + return scalar::cosine_int8_batch_distance; + } + if (metric_type == MetricType::kInnerProduct) { + return scalar::inner_product_int8_batch_distance; + } + } } if (quantize_type == QuantizeType::kUniform) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && @@ -126,7 +222,8 @@ QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, QuantizeType quantize_type, CpuArchType cpu_arch_type) { if (data_type == DataType::kInt8) { - if (quantize_type == QuantizeType::kDefault) { + if (quantize_type == QuantizeType::kDefault || + quantize_type == QuantizeType::kRecord) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { diff --git a/tests/turbo/turbo_fp16_quantizer_test.cc b/tests/turbo/turbo_fp16_quantizer_test.cc new file mode 100644 index 000000000..9e4ae3504 --- /dev/null +++ b/tests/turbo/turbo_fp16_quantizer_test.cc @@ -0,0 +1,258 @@ +// 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 "zvec/core/framework/index_factory.h" + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +// Helper: reference cosine distance between two raw fp32 vectors. +static float reference_cosine(const float *a, const float *b, size_t dim) { + float dot = 0.0f, na = 0.0f, nb = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom; +} + +TEST(Fp16Quantizer, General) { + std::mt19937 gen(15583); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t COUNT = 10000; + const size_t DIMENSION = 12; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Fp16Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + + auto holder = + std::make_shared>( + DIMENSION); + for (size_t i = 0; i < COUNT; ++i) { + zvec::ailego::NumericalVector vec(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + vec[j] = dist(gen); + } + holder->emplace(i + 1, vec); + } + EXPECT_EQ(COUNT, holder->count()); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, holder->data_type()); + + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::string quant_buffer; + std::string dequant_buffer; + + for (; iter->is_valid(); iter->next()) { + EXPECT_TRUE(iter->data()); + + IndexQueryMeta qmeta; + quant_buffer.clear(); + EXPECT_EQ(0, quantizer->quantize( + iter->data(), + IndexQueryMeta(holder->data_type(), holder->dimension()), + &quant_buffer, &qmeta)); + EXPECT_EQ(IndexMeta::DataType::DT_FP16, qmeta.data_type()); + EXPECT_EQ(holder->dimension(), qmeta.dimension()); + + dequant_buffer.clear(); + EXPECT_EQ( + 0, quantizer->dequantize(quant_buffer.data(), qmeta, &dequant_buffer)); + + const float *original_data = reinterpret_cast(iter->data()); + const float *dequantize_data = + reinterpret_cast(dequant_buffer.data()); + for (size_t i = 0; i < holder->dimension(); ++i) { + EXPECT_NEAR(original_data[i], dequantize_data[i], 1e-2); + } + } +} + +TEST(Fp16Quantizer, Score) { + std::mt19937 gen(42); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Fp16Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + + // Generate raw vectors and quantize them. + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + IndexQueryMeta ometa; + EXPECT_EQ(0, quantizer->quantize( + raw_vecs[i].data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, DIMENSION), + &quant_vecs[i], &ometa)); + } + + // --- calc_distance_dp_query (single) --- + for (size_t i = 1; i < COUNT; ++i) { + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-2) << "i=" << i; + } + + // --- calc_distance_dp_query_batch --- + { + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch( + dp_list.data(), static_cast(dp_list.size()), quant_vecs[0].data(), + results.data()); + + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[i + 1].data(), + raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(results[i], expected, 1e-2) << "i=" << i; + } + } + + // --- distance() + DistanceImpl (single + batch) --- + { + IndexQueryMeta qmeta(IndexMeta::MetaType::MT_DENSE, + IndexMeta::DataType::DT_FP16, + IndexMeta::UnitSizeof(IndexMeta::DataType::DT_FP16), + DIMENSION, 0, sizeof(float)); + auto dist_impl = quantizer->distance(quant_vecs[0].data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + + for (size_t i = 1; i < COUNT; ++i) { + float d = dist_impl(quant_vecs[i].data()); + float expected = + reference_cosine(raw_vecs[0].data(), raw_vecs[i].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-2) << "i=" << i; + } + + // Batch via DistanceImpl. + ASSERT_TRUE(dist_impl.batch_valid()); + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector batch_results(COUNT - 1); + dist_impl.batch(dp_list.data(), dp_list.size(), batch_results.data()); + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[0].data(), + raw_vecs[i + 1].data(), DIMENSION); + EXPECT_NEAR(batch_results[i], expected, 1e-2) << "i=" << i; + } + } + + // --- calc_distance_dp_dp (pairwise) --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_dp(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-2) << "i=" << i; + } + + // --- calc_distance_dp_query_unquantized --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_query_unquantized( + quant_vecs[i].data(), raw_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 1e-2) << "i=" << i; + } +} + +TEST(Fp16Quantizer, ScoreSquaredEuclidean) { + std::mt19937 gen(7); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("SquaredEuclidean", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Fp16Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + EXPECT_EQ(DIMENSION * sizeof(uint16_t), + quantizer->quantized_datapoint_vector_length()); + + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + quant_vecs[i].resize(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(raw_vecs[i].data(), &quant_vecs[i][0]); + } + + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch(dp_list.data(), + static_cast(dp_list.size()), + quant_vecs[0].data(), results.data()); + + for (size_t i = 1; i < COUNT; ++i) { + float expected = 0.0f; + for (size_t j = 0; j < DIMENSION; ++j) { + float diff = raw_vecs[i][j] - raw_vecs[0][j]; + expected += diff * diff; + } + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + EXPECT_NEAR(d, expected, 1e-2) << "i=" << i; + EXPECT_NEAR(results[i - 1], expected, 1e-2) << "i=" << i; + } +} diff --git a/tests/turbo/turbo_int4_quantizer_test.cc b/tests/turbo/turbo_int4_quantizer_test.cc new file mode 100644 index 000000000..7aa380f5f --- /dev/null +++ b/tests/turbo/turbo_int4_quantizer_test.cc @@ -0,0 +1,280 @@ +// 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 "zvec/core/framework/index_factory.h" + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +// Record tail size for int4: 4 floats. +static constexpr size_t kInt4TailSize = 16; + +// Helper: reference cosine distance between two raw fp32 vectors. +static float reference_cosine(const float *a, const float *b, size_t dim) { + float dot = 0.0f, na = 0.0f, nb = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom; +} + +TEST(Int4Quantizer, General) { + std::mt19937 gen(15583); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t COUNT = 10000; + const size_t DIMENSION = 12; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Int4Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + EXPECT_EQ(DIMENSION / 2 + kInt4TailSize + sizeof(float), + quantizer->quantized_datapoint_vector_length()); + + auto holder = + std::make_shared>( + DIMENSION); + for (size_t i = 0; i < COUNT; ++i) { + zvec::ailego::NumericalVector vec(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + vec[j] = dist(gen); + } + holder->emplace(i + 1, vec); + } + EXPECT_EQ(COUNT, holder->count()); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, holder->data_type()); + + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::string quant_buffer; + std::string dequant_buffer; + + for (; iter->is_valid(); iter->next()) { + EXPECT_TRUE(iter->data()); + + IndexQueryMeta qmeta; + quant_buffer.clear(); + EXPECT_EQ(0, quantizer->quantize( + iter->data(), + IndexQueryMeta(holder->data_type(), holder->dimension()), + &quant_buffer, &qmeta)); + EXPECT_EQ(IndexMeta::DataType::DT_INT4, qmeta.data_type()); + EXPECT_EQ(holder->dimension(), qmeta.dimension()); + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), + quant_buffer.size()); + + dequant_buffer.clear(); + EXPECT_EQ( + 0, quantizer->dequantize(quant_buffer.data(), qmeta, &dequant_buffer)); + + const float *original_data = reinterpret_cast(iter->data()); + const float *dequantize_data = + reinterpret_cast(dequant_buffer.data()); + for (size_t i = 0; i < holder->dimension(); ++i) { + EXPECT_NEAR(original_data[i], dequantize_data[i], 0.15); + } + } +} + +TEST(Int4Quantizer, OddDimensionRejected) { + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, 13); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Int4Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + EXPECT_NE(0, quantizer->init(meta, params)); +} + +TEST(Int4Quantizer, Score) { + std::mt19937 gen(42); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Int4Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + + // Generate raw vectors and quantize them. + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + IndexQueryMeta ometa; + EXPECT_EQ(0, quantizer->quantize( + raw_vecs[i].data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, DIMENSION), + &quant_vecs[i], &ometa)); + } + + // --- calc_distance_dp_query (single) --- + for (size_t i = 1; i < COUNT; ++i) { + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 0.1) << "i=" << i; + } + + // --- calc_distance_dp_query_batch --- + { + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch( + dp_list.data(), static_cast(dp_list.size()), quant_vecs[0].data(), + results.data()); + + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[i + 1].data(), + raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(results[i], expected, 0.1) << "i=" << i; + } + } + + // --- distance() + DistanceImpl (single + batch) --- + { + IndexQueryMeta qmeta; + qmeta.set_meta(IndexMeta::DataType::DT_INT4, DIMENSION, + static_cast(turbo::QuantizeType::kRecord), + kInt4TailSize + sizeof(float)); + auto dist_impl = quantizer->distance(quant_vecs[0].data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + + for (size_t i = 1; i < COUNT; ++i) { + float d = dist_impl(quant_vecs[i].data()); + float expected = + reference_cosine(raw_vecs[0].data(), raw_vecs[i].data(), DIMENSION); + EXPECT_NEAR(d, expected, 0.1) << "i=" << i; + } + + // Batch via DistanceImpl. + ASSERT_TRUE(dist_impl.batch_valid()); + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector batch_results(COUNT - 1); + dist_impl.batch(dp_list.data(), dp_list.size(), batch_results.data()); + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[0].data(), + raw_vecs[i + 1].data(), DIMENSION); + EXPECT_NEAR(batch_results[i], expected, 0.1) << "i=" << i; + } + } + + // --- calc_distance_dp_dp (pairwise) --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_dp(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 0.1) << "i=" << i; + } + + // --- calc_distance_dp_query_unquantized --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_query_unquantized( + quant_vecs[i].data(), raw_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 0.1) << "i=" << i; + } +} + +TEST(Int4Quantizer, ScoreSquaredEuclidean) { + std::mt19937 gen(7); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("SquaredEuclidean", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Int4Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + EXPECT_EQ(DIMENSION / 2 + kInt4TailSize, + quantizer->quantized_datapoint_vector_length()); + + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + quant_vecs[i].resize(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(raw_vecs[i].data(), &quant_vecs[i][0]); + } + + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch(dp_list.data(), + static_cast(dp_list.size()), + quant_vecs[0].data(), results.data()); + + for (size_t i = 1; i < COUNT; ++i) { + float expected = 0.0f; + for (size_t j = 0; j < DIMENSION; ++j) { + float diff = raw_vecs[i][j] - raw_vecs[0][j]; + expected += diff * diff; + } + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + EXPECT_NEAR(d, expected, 0.5) << "i=" << i; + EXPECT_NEAR(results[i - 1], expected, 0.5) << "i=" << i; + + float du = quantizer->calc_distance_dp_query_unquantized( + quant_vecs[i].data(), raw_vecs[0].data()); + EXPECT_NEAR(du, expected, 0.5) << "i=" << i; + } +} diff --git a/tests/turbo/turbo_int8_quantizer_test.cc b/tests/turbo/turbo_int8_quantizer_test.cc new file mode 100644 index 000000000..5127fb38d --- /dev/null +++ b/tests/turbo/turbo_int8_quantizer_test.cc @@ -0,0 +1,269 @@ +// 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 "zvec/core/framework/index_factory.h" + +using namespace zvec; +using namespace zvec::core; +using namespace zvec::ailego; + +// Record tail size for int8: 4 floats + 1 int32. +static constexpr size_t kInt8TailSize = 20; + +// Helper: reference cosine distance between two raw fp32 vectors. +static float reference_cosine(const float *a, const float *b, size_t dim) { + float dot = 0.0f, na = 0.0f, nb = 0.0f; + for (size_t i = 0; i < dim; ++i) { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + float denom = std::sqrt(na) * std::sqrt(nb); + return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom; +} + +TEST(Int8Quantizer, General) { + std::mt19937 gen(15583); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t COUNT = 10000; + const size_t DIMENSION = 12; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Int8Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + EXPECT_EQ(DIMENSION + kInt8TailSize + sizeof(float), + quantizer->quantized_datapoint_vector_length()); + + auto holder = + std::make_shared>( + DIMENSION); + for (size_t i = 0; i < COUNT; ++i) { + zvec::ailego::NumericalVector vec(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + vec[j] = dist(gen); + } + holder->emplace(i + 1, vec); + } + EXPECT_EQ(COUNT, holder->count()); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, holder->data_type()); + + ASSERT_EQ(0, quantizer->train(holder)); + + auto iter = holder->create_iterator(); + std::string quant_buffer; + std::string dequant_buffer; + + for (; iter->is_valid(); iter->next()) { + EXPECT_TRUE(iter->data()); + + IndexQueryMeta qmeta; + quant_buffer.clear(); + EXPECT_EQ(0, quantizer->quantize( + iter->data(), + IndexQueryMeta(holder->data_type(), holder->dimension()), + &quant_buffer, &qmeta)); + EXPECT_EQ(IndexMeta::DataType::DT_INT8, qmeta.data_type()); + EXPECT_EQ(holder->dimension(), qmeta.dimension()); + EXPECT_EQ(quantizer->quantized_datapoint_vector_length(), + quant_buffer.size()); + + dequant_buffer.clear(); + EXPECT_EQ( + 0, quantizer->dequantize(quant_buffer.data(), qmeta, &dequant_buffer)); + + const float *original_data = reinterpret_cast(iter->data()); + const float *dequantize_data = + reinterpret_cast(dequant_buffer.data()); + for (size_t i = 0; i < holder->dimension(); ++i) { + EXPECT_NEAR(original_data[i], dequantize_data[i], 1e-2); + } + } +} + +TEST(Int8Quantizer, Score) { + std::mt19937 gen(42); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("Cosine", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Int8Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + + // Generate raw vectors and quantize them. + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + IndexQueryMeta ometa; + EXPECT_EQ(0, quantizer->quantize( + raw_vecs[i].data(), + IndexQueryMeta(IndexMeta::DataType::DT_FP32, DIMENSION), + &quant_vecs[i], &ometa)); + } + + // --- calc_distance_dp_query (single) --- + for (size_t i = 1; i < COUNT; ++i) { + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 2e-2) << "i=" << i; + } + + // --- calc_distance_dp_query_batch --- + { + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch( + dp_list.data(), static_cast(dp_list.size()), quant_vecs[0].data(), + results.data()); + + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[i + 1].data(), + raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(results[i], expected, 2e-2) << "i=" << i; + } + } + + // --- distance() + DistanceImpl (single + batch fallback) --- + { + IndexQueryMeta qmeta; + qmeta.set_meta(IndexMeta::DataType::DT_INT8, DIMENSION, + static_cast(turbo::QuantizeType::kRecord), + kInt8TailSize + sizeof(float)); + auto dist_impl = quantizer->distance(quant_vecs[0].data(), qmeta); + ASSERT_TRUE(dist_impl.valid()); + + for (size_t i = 1; i < COUNT; ++i) { + float d = dist_impl(quant_vecs[i].data()); + float expected = + reference_cosine(raw_vecs[0].data(), raw_vecs[i].data(), DIMENSION); + EXPECT_NEAR(d, expected, 2e-2) << "i=" << i; + } + + // Batch via DistanceImpl (falls back to the single-pair kernel when the + // batch kernel requires query preprocessing). + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector batch_results(COUNT - 1); + dist_impl.batch(dp_list.data(), dp_list.size(), batch_results.data()); + for (size_t i = 0; i < dp_list.size(); ++i) { + float expected = reference_cosine(raw_vecs[0].data(), + raw_vecs[i + 1].data(), DIMENSION); + EXPECT_NEAR(batch_results[i], expected, 2e-2) << "i=" << i; + } + } + + // --- calc_distance_dp_dp (pairwise) --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_dp(quant_vecs[i].data(), + quant_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 2e-2) << "i=" << i; + } + + // --- calc_distance_dp_query_unquantized --- + for (size_t i = 1; i < 10; ++i) { + float d = quantizer->calc_distance_dp_query_unquantized( + quant_vecs[i].data(), raw_vecs[0].data()); + float expected = + reference_cosine(raw_vecs[i].data(), raw_vecs[0].data(), DIMENSION); + EXPECT_NEAR(d, expected, 2e-2) << "i=" << i; + } +} + +TEST(Int8Quantizer, ScoreSquaredEuclidean) { + std::mt19937 gen(7); + std::uniform_real_distribution dist(0.0, 1.0); + + const size_t DIMENSION = 12; + const size_t COUNT = 100; + + IndexMeta meta; + meta.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION); + meta.set_metric("SquaredEuclidean", 0, Params()); + + auto quantizer = IndexFactory::CreateQuantizer("Int8Quantizer"); + ASSERT_TRUE(quantizer); + zvec::ailego::Params params; + ASSERT_EQ(0, quantizer->init(meta, params)); + EXPECT_EQ(DIMENSION + kInt8TailSize, + quantizer->quantized_datapoint_vector_length()); + + std::vector> raw_vecs(COUNT); + std::vector quant_vecs(COUNT); + for (size_t i = 0; i < COUNT; ++i) { + raw_vecs[i].resize(DIMENSION); + for (size_t j = 0; j < DIMENSION; ++j) { + raw_vecs[i][j] = dist(gen); + } + quant_vecs[i].resize(quantizer->quantized_datapoint_vector_length()); + quantizer->quantize_data(raw_vecs[i].data(), &quant_vecs[i][0]); + } + + std::vector dp_list(COUNT - 1); + for (size_t i = 1; i < COUNT; ++i) { + dp_list[i - 1] = quant_vecs[i].data(); + } + std::vector results(COUNT - 1); + quantizer->calc_distance_dp_query_batch(dp_list.data(), + static_cast(dp_list.size()), + quant_vecs[0].data(), results.data()); + + for (size_t i = 1; i < COUNT; ++i) { + float expected = 0.0f; + for (size_t j = 0; j < DIMENSION; ++j) { + float diff = raw_vecs[i][j] - raw_vecs[0][j]; + expected += diff * diff; + } + float d = quantizer->calc_distance_dp_query(quant_vecs[i].data(), + quant_vecs[0].data()); + EXPECT_NEAR(d, expected, 5e-2) << "i=" << i; + EXPECT_NEAR(results[i - 1], expected, 5e-2) << "i=" << i; + + float du = quantizer->calc_distance_dp_query_unquantized( + quant_vecs[i].data(), raw_vecs[0].data()); + EXPECT_NEAR(du, expected, 5e-2) << "i=" << i; + } +} From 2eeeefd3b595767f4d5e23fc27199affa1273b63 Mon Sep 17 00:00:00 2001 From: Jianning Wang Date: Tue, 28 Jul 2026 15:29:24 +0800 Subject: [PATCH 2/4] refactor: new findKernel --- src/core/metric/quantized_integer_metric.cc | 12 +- src/include/zvec/turbo/turbo.h | 20 +- .../fp16_quantizer/fp16_quantizer.cc | 21 +- .../quantizer/fp16_quantizer/fp16_quantizer.h | 4 +- .../fp32_quantizer/fp32_quantizer.cc | 21 +- .../quantizer/fp32_quantizer/fp32_quantizer.h | 4 +- .../int4_quantizer/int4_quantizer.cc | 10 +- .../quantizer/int4_quantizer/int4_quantizer.h | 4 +- .../int8_quantizer/int8_quantizer.cc | 13 +- .../quantizer/int8_quantizer/int8_quantizer.h | 4 +- src/turbo/quantizer/quantizer.h | 6 +- src/turbo/turbo.cc | 340 +++++++++--------- 12 files changed, 221 insertions(+), 238 deletions(-) diff --git a/src/core/metric/quantized_integer_metric.cc b/src/core/metric/quantized_integer_metric.cc index 1d46f7749..0fd8a2f3c 100644 --- a/src/core/metric/quantized_integer_metric.cc +++ b/src/core/metric/quantized_integer_metric.cc @@ -98,7 +98,7 @@ class QuantizedIntegerMetric : public IndexMetric { if (meta_.data_type() == IndexMeta::DataType::DT_INT8) { auto turbo_ret = turbo::get_distance_func( turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8, - turbo::QuantizeType::kDefault); + turbo::QuantizeType::kRecord, turbo::CpuArchType::kAVX512VNNI); if (turbo_ret && m == 1 && n == 1) { return turbo_ret; } @@ -139,7 +139,7 @@ class QuantizedIntegerMetric : public IndexMetric { if (meta_.data_type() == IndexMeta::DataType::DT_INT8) { auto turbo_ret = turbo::get_distance_func( turbo::MetricType::kCosine, turbo::DataType::kInt8, - turbo::QuantizeType::kDefault); + turbo::QuantizeType::kRecord, turbo::CpuArchType::kAVX512VNNI); if (turbo_ret) { return turbo_ret; } @@ -160,7 +160,7 @@ class QuantizedIntegerMetric : public IndexMetric { if (meta_.data_type() == IndexMeta::DataType::DT_INT8) { auto turbo_ret = turbo::get_batch_distance_func( turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8, - turbo::QuantizeType::kDefault); + turbo::QuantizeType::kRecord, turbo::CpuArchType::kAVX512VNNI); if (turbo_ret) { return turbo_ret; } @@ -215,7 +215,7 @@ class QuantizedIntegerMetric : public IndexMetric { if (meta_.data_type() == IndexMeta::DataType::DT_INT8) { auto turbo_ret = turbo::get_batch_distance_func( turbo::MetricType::kCosine, turbo::DataType::kInt8, - turbo::QuantizeType::kDefault); + turbo::QuantizeType::kRecord, turbo::CpuArchType::kAVX512VNNI); if (turbo_ret) { return turbo_ret; } @@ -300,7 +300,7 @@ class QuantizedIntegerMetric : public IndexMetric { meta_.data_type() == IndexMeta::DataType::DT_INT8) { auto turbo_ret = turbo::get_query_preprocess_func( turbo::MetricType::kCosine, turbo::DataType::kInt8, - turbo::QuantizeType::kDefault); + turbo::QuantizeType::kRecord, turbo::CpuArchType::kAVX512VNNI); if (turbo_ret) { return turbo_ret; } @@ -310,7 +310,7 @@ class QuantizedIntegerMetric : public IndexMetric { meta_.data_type() == IndexMeta::DataType::DT_INT8) { auto turbo_ret = turbo::get_query_preprocess_func( turbo::MetricType::kSquaredEuclidean, turbo::DataType::kInt8, - turbo::QuantizeType::kDefault); + turbo::QuantizeType::kRecord, turbo::CpuArchType::kAVX512VNNI); if (turbo_ret) { return turbo_ret; } diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index 5baa8c418..51a8a25a6 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -88,7 +88,10 @@ enum class DataType { }; enum class QuantizeType { - kDefault, + //! Deprecated: no dispatch row serves kDefault anymore; request the + //! explicit quantize type (kFp32, kFp16, kRecord, ...) instead. The + //! enumerator is kept (value 0) for serialized-header compatibility. + kDefault [[deprecated("request an explicit QuantizeType instead")]], kUniform, kRecord, kFp16, @@ -129,6 +132,21 @@ ZVEC_TURBO_API QueryPreprocessFunc get_query_preprocess_func( MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type = CpuArchType::kAuto); +// All kernels of a single dispatched kernel family. `preprocess` is non-null +// when the batch kernel requires the query to be preprocessed first (e.g. +// the AVX512-VNNI int8 kernels expect a +128 uint8-shifted query). +struct DistanceKernels { + DistanceFunc dist{}; + BatchDistanceFunc batch{}; + QueryPreprocessFunc preprocess = nullptr; +}; + +// Aggregate lookup: resolves dist/batch/preprocess in one pass so callers +// cannot pair functions from different kernel families. +DistanceKernels get_distance_kernels( + MetricType metric_type, DataType data_type, QuantizeType quantize_type, + CpuArchType cpu_arch_type = CpuArchType::kAuto); + // Returns the SIMD kernel for the uniform quantizer on the current CPU for // the given output data_type, or nullptr if no SIMD implementation is // available (callers must keep a scalar fallback). This is a diff --git a/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc index 7ff6517ef..0de5b06ef 100644 --- a/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc +++ b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.cc @@ -38,12 +38,11 @@ int Fp16Quantizer::init(const IndexMeta &meta, } // Cache the distance dispatch for the new Quantizer interface. - dp_query_func_ = - get_distance_func(metric_from_name(metric_name), DataType::kFp16, - QuantizeType::kFp16, CpuArchType::kAuto); - dp_query_batch_func_ = - get_batch_distance_func(metric_from_name(metric_name), DataType::kFp16, - QuantizeType::kFp16, CpuArchType::kAuto); + auto kernels = + get_distance_kernels(metric_from_name(metric_name), DataType::kFp16, + QuantizeType::kFp16, CpuArchType::kAuto); + dp_query_func_ = std::move(kernels.dist); + dp_query_batch_func_ = std::move(kernels.batch); return 0; } @@ -116,19 +115,15 @@ int Fp16Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, DistanceImpl Fp16Quantizer::distance(const void *query, const IndexQueryMeta &qmeta) const { - auto metric = metric_from_name(meta_.metric_name()); - auto func = get_distance_func(metric, DataType::kFp16, QuantizeType::kFp16, - CpuArchType::kAuto); - if (!func) { + // Reuse the dispatch cached at init(). + if (!dp_query_func_) { return DistanceImpl{}; } - auto batch_func = get_batch_distance_func( - metric, DataType::kFp16, QuantizeType::kFp16, CpuArchType::kAuto); // The query is assumed to be already quantized — copy it directly. std::string quantized_query(static_cast(query), qmeta.element_size()); - return DistanceImpl(std::move(func), std::move(batch_func), + return DistanceImpl(dp_query_func_, dp_query_batch_func_, std::move(quantized_query), original_dim_); } diff --git a/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h index 6b87b86cd..4285e89ac 100644 --- a/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h +++ b/src/turbo/quantizer/fp16_quantizer/fp16_quantizer.h @@ -27,9 +27,7 @@ using namespace zvec::core; class Fp16Quantizer : public Quantizer { public: - Fp16Quantizer() { - type_ = QuantizeType::kFp16; - } + Fp16Quantizer() : Quantizer(QuantizeType::kFp16) {} virtual ~Fp16Quantizer() {} diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index 839137e6d..c473f6df6 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -39,12 +39,11 @@ int Fp32Quantizer::init(const IndexMeta &meta, } // Cache the distance dispatch for the new Quantizer interface. - dp_query_func_ = - get_distance_func(metric_from_name(metric_name), DataType::kFp32, - QuantizeType::kDefault, CpuArchType::kAuto); - dp_query_batch_func_ = - get_batch_distance_func(metric_from_name(metric_name), DataType::kFp32, - QuantizeType::kDefault, CpuArchType::kAuto); + auto kernels = + get_distance_kernels(metric_from_name(metric_name), DataType::kFp32, + QuantizeType::kFp32, CpuArchType::kAuto); + dp_query_func_ = std::move(kernels.dist); + dp_query_batch_func_ = std::move(kernels.batch); return 0; } @@ -112,19 +111,15 @@ int Fp32Quantizer::dequantize(const void *in, const IndexQueryMeta &qmeta, DistanceImpl Fp32Quantizer::distance(const void *query, const IndexQueryMeta &qmeta) const { - auto metric = metric_from_name(meta_.metric_name()); - auto func = get_distance_func(metric, DataType::kFp32, QuantizeType::kDefault, - CpuArchType::kAuto); - if (!func) { + // Reuse the dispatch cached at init(). + if (!dp_query_func_) { return DistanceImpl{}; } - auto batch_func = get_batch_distance_func( - metric, DataType::kFp32, QuantizeType::kDefault, CpuArchType::kAuto); // The query is assumed to be already quantized — copy it directly. std::string quantized_query(static_cast(query), qmeta.element_size()); - return DistanceImpl(std::move(func), std::move(batch_func), + return DistanceImpl(dp_query_func_, dp_query_batch_func_, std::move(quantized_query), original_dim_); } diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 885390639..2d5d5906d 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -27,9 +27,7 @@ using namespace zvec::core; class Fp32Quantizer : public Quantizer { public: - Fp32Quantizer() { - type_ = QuantizeType::kFp32; - } + Fp32Quantizer() : Quantizer(QuantizeType::kFp32) {} virtual ~Fp32Quantizer() {} diff --git a/src/turbo/quantizer/int4_quantizer/int4_quantizer.cc b/src/turbo/quantizer/int4_quantizer/int4_quantizer.cc index 5bb76d47a..dcee5b228 100644 --- a/src/turbo/quantizer/int4_quantizer/int4_quantizer.cc +++ b/src/turbo/quantizer/int4_quantizer/int4_quantizer.cc @@ -56,11 +56,11 @@ int Int4Quantizer::init(const IndexMeta &meta, dist_dim_ = static_cast(original_dim_) + 2 * extra_meta_size_; // Cache the distance dispatch for the new Quantizer interface. - auto metric = metric_from_name(metric_name); - dp_query_func_ = get_distance_func(metric, DataType::kInt4, - QuantizeType::kRecord, CpuArchType::kAuto); - dp_query_batch_func_ = get_batch_distance_func( - metric, DataType::kInt4, QuantizeType::kRecord, CpuArchType::kAuto); + auto kernels = + get_distance_kernels(metric_from_name(metric_name), DataType::kInt4, + QuantizeType::kRecord, CpuArchType::kAuto); + dp_query_func_ = std::move(kernels.dist); + dp_query_batch_func_ = std::move(kernels.batch); return 0; } diff --git a/src/turbo/quantizer/int4_quantizer/int4_quantizer.h b/src/turbo/quantizer/int4_quantizer/int4_quantizer.h index 435e8ebcd..93c74cd8a 100644 --- a/src/turbo/quantizer/int4_quantizer/int4_quantizer.h +++ b/src/turbo/quantizer/int4_quantizer/int4_quantizer.h @@ -38,9 +38,7 @@ using namespace zvec::core; */ class Int4Quantizer : public Quantizer { public: - Int4Quantizer() { - type_ = QuantizeType::kRecord; - } + Int4Quantizer() : Quantizer(QuantizeType::kRecord) {} virtual ~Int4Quantizer() {} diff --git a/src/turbo/quantizer/int8_quantizer/int8_quantizer.cc b/src/turbo/quantizer/int8_quantizer/int8_quantizer.cc index b6c8bfeb4..78dd8efec 100644 --- a/src/turbo/quantizer/int8_quantizer/int8_quantizer.cc +++ b/src/turbo/quantizer/int8_quantizer/int8_quantizer.cc @@ -50,13 +50,12 @@ int Int8Quantizer::init(const IndexMeta &meta, dist_dim_ = static_cast(original_dim_) + extra_meta_size_; // Cache the distance dispatch for the new Quantizer interface. - auto metric = metric_from_name(metric_name); - dp_query_func_ = get_distance_func(metric, DataType::kInt8, - QuantizeType::kRecord, CpuArchType::kAuto); - dp_query_batch_func_ = get_batch_distance_func( - metric, DataType::kInt8, QuantizeType::kRecord, CpuArchType::kAuto); - dp_query_preprocess_func_ = get_query_preprocess_func( - metric, DataType::kInt8, QuantizeType::kRecord, CpuArchType::kAuto); + auto kernels = + get_distance_kernels(metric_from_name(metric_name), DataType::kInt8, + QuantizeType::kRecord, CpuArchType::kAuto); + dp_query_func_ = std::move(kernels.dist); + dp_query_batch_func_ = std::move(kernels.batch); + dp_query_preprocess_func_ = kernels.preprocess; return 0; } diff --git a/src/turbo/quantizer/int8_quantizer/int8_quantizer.h b/src/turbo/quantizer/int8_quantizer/int8_quantizer.h index d8003c5d2..874c5e71b 100644 --- a/src/turbo/quantizer/int8_quantizer/int8_quantizer.h +++ b/src/turbo/quantizer/int8_quantizer/int8_quantizer.h @@ -35,9 +35,7 @@ using namespace zvec::core; */ class Int8Quantizer : public Quantizer { public: - Int8Quantizer() { - type_ = QuantizeType::kRecord; - } + Int8Quantizer() : Quantizer(QuantizeType::kRecord) {} virtual ~Int8Quantizer() {} diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 77572a6a8..d6318800d 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -47,7 +47,6 @@ class Quantizer { public: typedef std::shared_ptr Pointer; - Quantizer() {} virtual ~Quantizer() {} //! Initialize quantizer with index metadata and parameters @@ -148,6 +147,9 @@ class Quantizer { } protected: + //! Subclasses must declare which QuantizeType they implement. + explicit Quantizer(QuantizeType type) : type_(type) {} + //! Map a metric name (e.g. "SquaredEuclidean", "Cosine", //! "InnerProduct", "MipsSquaredEuclidean") to its MetricType. static MetricType metric_from_name(const std::string &name) { @@ -166,7 +168,7 @@ class Quantizer { return MetricType::kUnknown; } - QuantizeType type_{QuantizeType::kDefault}; + QuantizeType type_; uint32_t extra_meta_size_{0}; }; diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 1b4330890..c75fc90e7 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -44,198 +44,180 @@ static bool IsArchMatch(CpuArchType actual, CpuArchType target) { return actual == CpuArchType::kAuto || actual == target; } -DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, - QuantizeType quantize_type, - CpuArchType cpu_arch_type) { - if (data_type == DataType::kFp32) { - if (quantize_type == QuantizeType::kDefault || - quantize_type == QuantizeType::kFp32) { - if (metric_type == MetricType::kCosine) { - return scalar::cosine_fp32_distance; - } - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_fp32_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_fp32_distance; - } - } - return nullptr; +// Single place that maps a CpuArchType to its runtime CPU-feature gate. +static bool CpuSupports(CpuArchType arch) { + const auto &flags = zvec::ailego::internal::CpuFeatures::static_flags_; + switch (arch) { + case CpuArchType::kScalar: + return true; + case CpuArchType::kSSE: + return flags.SSE2; + case CpuArchType::kAVX2: + return flags.AVX2; + case CpuArchType::kAVX512: + return flags.AVX512F; + case CpuArchType::kAVX512VNNI: + return flags.AVX512_VNNI; + case CpuArchType::kNEON: + return flags.NEON; + default: + return false; } - if (data_type == DataType::kFp16) { - if (quantize_type == QuantizeType::kDefault || - quantize_type == QuantizeType::kFp16) { - if (metric_type == MetricType::kCosine) { - return scalar::cosine_fp16_distance; - } - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_fp16_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_fp16_distance; - } - } - return nullptr; - } - if (data_type == DataType::kInt4) { - if (quantize_type == QuantizeType::kRecord) { - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_int4_distance; - } - if (metric_type == MetricType::kCosine) { - return scalar::cosine_int4_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_int4_distance; - } - } - return nullptr; - } - if (data_type == DataType::kInt8) { - if (quantize_type == QuantizeType::kDefault || - quantize_type == QuantizeType::kRecord) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { - if (metric_type == MetricType::kSquaredEuclidean) { - return avx512_vnni::squared_euclidean_int8_distance; - } - if (metric_type == MetricType::kCosine) { - return avx512_vnni::cosine_int8_distance; - } - } - // Scalar fallbacks are only exposed for the kRecord quantize type to - // keep the historical kDefault dispatch behavior unchanged. - if (quantize_type == QuantizeType::kRecord) { - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_int8_distance; - } - if (metric_type == MetricType::kCosine) { - return scalar::cosine_int8_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_int8_distance; - } - } +} + +namespace { + +// Raw kernel signatures (function pointers, so the registry can be a +// constexpr-friendly static table; they convert to the std::function-based +// public typedefs on return). +using RawDistanceFn = void (*)(const void *, const void *, size_t, float *); +using RawBatchDistanceFn = void (*)(const void *const *, const void *, size_t, + size_t, float *); + +constexpr uint32_t QBit(QuantizeType q) { + return 1u << static_cast(q); +} + +//! One row = one kernel family: all functions that must be used together +//! for a given (metric, data type) combination on a given ISA. +struct KernelSet { + MetricType metric; + DataType dtype; + uint32_t quantize_mask; //!< QuantizeTypes served by this row (QBit mask) + CpuArchType arch; //!< kScalar rows are the universal fallback + RawDistanceFn dist; + RawBatchDistanceFn batch; + QueryPreprocessFunc preprocess; //!< non-null: batch needs preprocessing +}; + +constexpr uint32_t kQmFp32 = QBit(QuantizeType::kFp32); +constexpr uint32_t kQmFp16 = QBit(QuantizeType::kFp16); +constexpr uint32_t kQmRecord = QBit(QuantizeType::kRecord); +constexpr uint32_t kQmUniform = QBit(QuantizeType::kUniform); + +// Dispatch registry. Row order encodes priority: SIMD rows must precede +// their scalar fallbacks. +constexpr KernelSet kKernelTable[] = { + // --- fp32 (scalar) --- + {MetricType::kCosine, DataType::kFp32, kQmFp32, CpuArchType::kScalar, + scalar::cosine_fp32_distance, scalar::cosine_fp32_batch_distance, nullptr}, + {MetricType::kSquaredEuclidean, DataType::kFp32, kQmFp32, + CpuArchType::kScalar, scalar::squared_euclidean_fp32_distance, + scalar::squared_euclidean_fp32_batch_distance, nullptr}, + {MetricType::kInnerProduct, DataType::kFp32, kQmFp32, CpuArchType::kScalar, + scalar::inner_product_fp32_distance, + scalar::inner_product_fp32_batch_distance, nullptr}, + + // --- fp16 (scalar) --- + {MetricType::kCosine, DataType::kFp16, kQmFp16, CpuArchType::kScalar, + scalar::cosine_fp16_distance, scalar::cosine_fp16_batch_distance, nullptr}, + {MetricType::kSquaredEuclidean, DataType::kFp16, kQmFp16, + CpuArchType::kScalar, scalar::squared_euclidean_fp16_distance, + scalar::squared_euclidean_fp16_batch_distance, nullptr}, + {MetricType::kInnerProduct, DataType::kFp16, kQmFp16, CpuArchType::kScalar, + scalar::inner_product_fp16_distance, + scalar::inner_product_fp16_batch_distance, nullptr}, + + // --- record-quantized int8 (AVX512-VNNI, then scalar fallback) --- + {MetricType::kSquaredEuclidean, DataType::kInt8, kQmRecord, + CpuArchType::kAVX512VNNI, avx512_vnni::squared_euclidean_int8_distance, + avx512_vnni::squared_euclidean_int8_batch_distance, + avx512_vnni::squared_euclidean_int8_query_preprocess}, + {MetricType::kCosine, DataType::kInt8, kQmRecord, CpuArchType::kAVX512VNNI, + avx512_vnni::cosine_int8_distance, avx512_vnni::cosine_int8_batch_distance, + avx512_vnni::cosine_int8_query_preprocess}, + {MetricType::kSquaredEuclidean, DataType::kInt8, kQmRecord, + CpuArchType::kScalar, scalar::squared_euclidean_int8_distance, + scalar::squared_euclidean_int8_batch_distance, nullptr}, + {MetricType::kCosine, DataType::kInt8, kQmRecord, CpuArchType::kScalar, + scalar::cosine_int8_distance, scalar::cosine_int8_batch_distance, nullptr}, + {MetricType::kInnerProduct, DataType::kInt8, kQmRecord, + CpuArchType::kScalar, scalar::inner_product_int8_distance, + scalar::inner_product_int8_batch_distance, nullptr}, + + // --- uniform-quantized int8 (AVX512-VNNI only) --- + {MetricType::kSquaredEuclidean, DataType::kInt8, kQmUniform, + CpuArchType::kAVX512VNNI, + avx512_vnni::uniform_squared_euclidean_int8_distance, + avx512_vnni::uniform_squared_euclidean_int8_batch_distance, nullptr}, + + // --- record-quantized int4 (scalar) --- + {MetricType::kSquaredEuclidean, DataType::kInt4, kQmRecord, + CpuArchType::kScalar, scalar::squared_euclidean_int4_distance, + scalar::squared_euclidean_int4_batch_distance, nullptr}, + {MetricType::kCosine, DataType::kInt4, kQmRecord, CpuArchType::kScalar, + scalar::cosine_int4_distance, scalar::cosine_int4_batch_distance, nullptr}, + {MetricType::kInnerProduct, DataType::kInt4, kQmRecord, + CpuArchType::kScalar, scalar::inner_product_int4_distance, + scalar::inner_product_int4_batch_distance, nullptr}, +}; + +// Returns the first (highest-priority) matching kernel row, or nullptr. +// Scalar rows are the fallback for auto dispatch: they match for kAuto and +// kScalar requests. An explicit SIMD arch request yields nullptr when that +// ISA is unavailable, so callers can keep their own (possibly SIMD-enabled) +// fallback paths instead of silently degrading to turbo scalar kernels. +const KernelSet *FindKernel(MetricType metric_type, DataType data_type, + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + for (const auto &k : kKernelTable) { + if (k.metric != metric_type || k.dtype != data_type || + (k.quantize_mask & QBit(quantize_type)) == 0) { + continue; } - if (quantize_type == QuantizeType::kUniform) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { - if (metric_type == MetricType::kSquaredEuclidean) { - return avx512_vnni::uniform_squared_euclidean_int8_distance; - } - } + if (IsArchMatch(cpu_arch_type, k.arch) && + (k.arch == CpuArchType::kScalar || CpuSupports(k.arch))) { + return &k; } } return nullptr; } -BatchDistanceFunc get_batch_distance_func(MetricType metric_type, - DataType data_type, - QuantizeType quantize_type, - CpuArchType cpu_arch_type) { - if (data_type == DataType::kFp32) { - if (quantize_type == QuantizeType::kDefault || - quantize_type == QuantizeType::kFp32) { - if (metric_type == MetricType::kCosine) { - return scalar::cosine_fp32_batch_distance; - } - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_fp32_batch_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_fp32_batch_distance; - } - } - return nullptr; +} // namespace + +DistanceKernels get_distance_kernels(MetricType metric_type, DataType data_type, + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + const KernelSet *k = + FindKernel(metric_type, data_type, quantize_type, cpu_arch_type); + if (!k) { + return DistanceKernels{}; } - if (data_type == DataType::kFp16) { - if (quantize_type == QuantizeType::kDefault || - quantize_type == QuantizeType::kFp16) { - if (metric_type == MetricType::kCosine) { - return scalar::cosine_fp16_batch_distance; - } - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_fp16_batch_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_fp16_batch_distance; - } - } - return nullptr; + DistanceKernels kernels; + if (k->dist) { + kernels.dist = k->dist; } - if (data_type == DataType::kInt4) { - if (quantize_type == QuantizeType::kRecord) { - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_int4_batch_distance; - } - if (metric_type == MetricType::kCosine) { - return scalar::cosine_int4_batch_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_int4_batch_distance; - } - } - return nullptr; - } - if (data_type == DataType::kInt8) { - if (quantize_type == QuantizeType::kDefault || - quantize_type == QuantizeType::kRecord) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { - if (metric_type == MetricType::kSquaredEuclidean) { - return avx512_vnni::squared_euclidean_int8_batch_distance; - } - if (metric_type == MetricType::kCosine) { - return avx512_vnni::cosine_int8_batch_distance; - } - } - // Scalar fallbacks are only exposed for the kRecord quantize type to - // keep the historical kDefault dispatch behavior unchanged. The scalar - // batch kernels take a plain int8 query (no preprocessing required). - if (quantize_type == QuantizeType::kRecord) { - if (metric_type == MetricType::kSquaredEuclidean) { - return scalar::squared_euclidean_int8_batch_distance; - } - if (metric_type == MetricType::kCosine) { - return scalar::cosine_int8_batch_distance; - } - if (metric_type == MetricType::kInnerProduct) { - return scalar::inner_product_int8_batch_distance; - } - } - } - if (quantize_type == QuantizeType::kUniform) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { - if (metric_type == MetricType::kSquaredEuclidean) { - return avx512_vnni::uniform_squared_euclidean_int8_batch_distance; - } - } - } + if (k->batch) { + kernels.batch = k->batch; } + kernels.preprocess = k->preprocess; + return kernels; +} - return nullptr; +DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + return get_distance_kernels(metric_type, data_type, quantize_type, + cpu_arch_type) + .dist; +} + +BatchDistanceFunc get_batch_distance_func(MetricType metric_type, + DataType data_type, + QuantizeType quantize_type, + CpuArchType cpu_arch_type) { + return get_distance_kernels(metric_type, data_type, quantize_type, + cpu_arch_type) + .batch; } QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type, DataType data_type, QuantizeType quantize_type, CpuArchType cpu_arch_type) { - if (data_type == DataType::kInt8) { - if (quantize_type == QuantizeType::kDefault || - quantize_type == QuantizeType::kRecord) { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { - if (metric_type == MetricType::kSquaredEuclidean) { - return avx512_vnni::squared_euclidean_int8_query_preprocess; - } - if (metric_type == MetricType::kCosine) { - return avx512_vnni::cosine_int8_query_preprocess; - } - } - } - } - return nullptr; + return get_distance_kernels(metric_type, data_type, quantize_type, + cpu_arch_type) + .preprocess; } UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { @@ -254,19 +236,19 @@ RotatorKernels get_rotator_kernels(RotateType rotate_type, CpuArchType cpu_arch_type) { switch (rotate_type) { case RotateType::kFht: { - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + if (CpuSupports(CpuArchType::kAVX512) && IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { return {avx512::fht_rotate_avx512, avx512::fht_unrotate_avx512}; } - if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + if (CpuSupports(CpuArchType::kAVX2) && IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { return {avx2::fht_rotate_avx2, avx2::fht_unrotate_avx2}; } - if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && + if (CpuSupports(CpuArchType::kSSE) && IsArchMatch(cpu_arch_type, CpuArchType::kSSE)) { return {sse::fht_rotate_sse, sse::fht_unrotate_sse}; } - if (zvec::ailego::internal::CpuFeatures::static_flags_.NEON && + if (CpuSupports(CpuArchType::kNEON) && IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { return {neon::fht_rotate_neon, neon::fht_unrotate_neon}; } From 1564ff6faaf3593c5bb8e5df0aa06f20620bf338 Mon Sep 17 00:00:00 2001 From: Jianning Wang Date: Tue, 28 Jul 2026 15:35:28 +0800 Subject: [PATCH 3/4] fix: don't use quantize mask --- src/turbo/turbo.cc | 64 ++++++++++++++++++++++------------------------ 1 file changed, 30 insertions(+), 34 deletions(-) diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index c75fc90e7..c4394d0df 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -74,80 +74,76 @@ using RawDistanceFn = void (*)(const void *, const void *, size_t, float *); using RawBatchDistanceFn = void (*)(const void *const *, const void *, size_t, size_t, float *); -constexpr uint32_t QBit(QuantizeType q) { - return 1u << static_cast(q); -} - //! One row = one kernel family: all functions that must be used together //! for a given (metric, data type) combination on a given ISA. struct KernelSet { MetricType metric; DataType dtype; - uint32_t quantize_mask; //!< QuantizeTypes served by this row (QBit mask) - CpuArchType arch; //!< kScalar rows are the universal fallback + QuantizeType quantize; //!< QuantizeType served by this row + CpuArchType arch; //!< kScalar rows are the universal fallback RawDistanceFn dist; RawBatchDistanceFn batch; QueryPreprocessFunc preprocess; //!< non-null: batch needs preprocessing }; -constexpr uint32_t kQmFp32 = QBit(QuantizeType::kFp32); -constexpr uint32_t kQmFp16 = QBit(QuantizeType::kFp16); -constexpr uint32_t kQmRecord = QBit(QuantizeType::kRecord); -constexpr uint32_t kQmUniform = QBit(QuantizeType::kUniform); - // Dispatch registry. Row order encodes priority: SIMD rows must precede // their scalar fallbacks. constexpr KernelSet kKernelTable[] = { // --- fp32 (scalar) --- - {MetricType::kCosine, DataType::kFp32, kQmFp32, CpuArchType::kScalar, - scalar::cosine_fp32_distance, scalar::cosine_fp32_batch_distance, nullptr}, - {MetricType::kSquaredEuclidean, DataType::kFp32, kQmFp32, + {MetricType::kCosine, DataType::kFp32, QuantizeType::kFp32, + CpuArchType::kScalar, scalar::cosine_fp32_distance, + scalar::cosine_fp32_batch_distance, nullptr}, + {MetricType::kSquaredEuclidean, DataType::kFp32, QuantizeType::kFp32, CpuArchType::kScalar, scalar::squared_euclidean_fp32_distance, scalar::squared_euclidean_fp32_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kFp32, kQmFp32, CpuArchType::kScalar, - scalar::inner_product_fp32_distance, + {MetricType::kInnerProduct, DataType::kFp32, QuantizeType::kFp32, + CpuArchType::kScalar, scalar::inner_product_fp32_distance, scalar::inner_product_fp32_batch_distance, nullptr}, // --- fp16 (scalar) --- - {MetricType::kCosine, DataType::kFp16, kQmFp16, CpuArchType::kScalar, - scalar::cosine_fp16_distance, scalar::cosine_fp16_batch_distance, nullptr}, - {MetricType::kSquaredEuclidean, DataType::kFp16, kQmFp16, + {MetricType::kCosine, DataType::kFp16, QuantizeType::kFp16, + CpuArchType::kScalar, scalar::cosine_fp16_distance, + scalar::cosine_fp16_batch_distance, nullptr}, + {MetricType::kSquaredEuclidean, DataType::kFp16, QuantizeType::kFp16, CpuArchType::kScalar, scalar::squared_euclidean_fp16_distance, scalar::squared_euclidean_fp16_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kFp16, kQmFp16, CpuArchType::kScalar, - scalar::inner_product_fp16_distance, + {MetricType::kInnerProduct, DataType::kFp16, QuantizeType::kFp16, + CpuArchType::kScalar, scalar::inner_product_fp16_distance, scalar::inner_product_fp16_batch_distance, nullptr}, // --- record-quantized int8 (AVX512-VNNI, then scalar fallback) --- - {MetricType::kSquaredEuclidean, DataType::kInt8, kQmRecord, + {MetricType::kSquaredEuclidean, DataType::kInt8, QuantizeType::kRecord, CpuArchType::kAVX512VNNI, avx512_vnni::squared_euclidean_int8_distance, avx512_vnni::squared_euclidean_int8_batch_distance, avx512_vnni::squared_euclidean_int8_query_preprocess}, - {MetricType::kCosine, DataType::kInt8, kQmRecord, CpuArchType::kAVX512VNNI, - avx512_vnni::cosine_int8_distance, avx512_vnni::cosine_int8_batch_distance, + {MetricType::kCosine, DataType::kInt8, QuantizeType::kRecord, + CpuArchType::kAVX512VNNI, avx512_vnni::cosine_int8_distance, + avx512_vnni::cosine_int8_batch_distance, avx512_vnni::cosine_int8_query_preprocess}, - {MetricType::kSquaredEuclidean, DataType::kInt8, kQmRecord, + {MetricType::kSquaredEuclidean, DataType::kInt8, QuantizeType::kRecord, CpuArchType::kScalar, scalar::squared_euclidean_int8_distance, scalar::squared_euclidean_int8_batch_distance, nullptr}, - {MetricType::kCosine, DataType::kInt8, kQmRecord, CpuArchType::kScalar, - scalar::cosine_int8_distance, scalar::cosine_int8_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kInt8, kQmRecord, + {MetricType::kCosine, DataType::kInt8, QuantizeType::kRecord, + CpuArchType::kScalar, scalar::cosine_int8_distance, + scalar::cosine_int8_batch_distance, nullptr}, + {MetricType::kInnerProduct, DataType::kInt8, QuantizeType::kRecord, CpuArchType::kScalar, scalar::inner_product_int8_distance, scalar::inner_product_int8_batch_distance, nullptr}, // --- uniform-quantized int8 (AVX512-VNNI only) --- - {MetricType::kSquaredEuclidean, DataType::kInt8, kQmUniform, + {MetricType::kSquaredEuclidean, DataType::kInt8, QuantizeType::kUniform, CpuArchType::kAVX512VNNI, avx512_vnni::uniform_squared_euclidean_int8_distance, avx512_vnni::uniform_squared_euclidean_int8_batch_distance, nullptr}, // --- record-quantized int4 (scalar) --- - {MetricType::kSquaredEuclidean, DataType::kInt4, kQmRecord, + {MetricType::kSquaredEuclidean, DataType::kInt4, QuantizeType::kRecord, CpuArchType::kScalar, scalar::squared_euclidean_int4_distance, scalar::squared_euclidean_int4_batch_distance, nullptr}, - {MetricType::kCosine, DataType::kInt4, kQmRecord, CpuArchType::kScalar, - scalar::cosine_int4_distance, scalar::cosine_int4_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kInt4, kQmRecord, + {MetricType::kCosine, DataType::kInt4, QuantizeType::kRecord, + CpuArchType::kScalar, scalar::cosine_int4_distance, + scalar::cosine_int4_batch_distance, nullptr}, + {MetricType::kInnerProduct, DataType::kInt4, QuantizeType::kRecord, CpuArchType::kScalar, scalar::inner_product_int4_distance, scalar::inner_product_int4_batch_distance, nullptr}, }; @@ -162,7 +158,7 @@ const KernelSet *FindKernel(MetricType metric_type, DataType data_type, CpuArchType cpu_arch_type) { for (const auto &k : kKernelTable) { if (k.metric != metric_type || k.dtype != data_type || - (k.quantize_mask & QBit(quantize_type)) == 0) { + k.quantize != quantize_type) { continue; } if (IsArchMatch(cpu_arch_type, k.arch) && From 943aab10d349e668dfdc458dcbff2314fb519488 Mon Sep 17 00:00:00 2001 From: Jianning Wang Date: Tue, 28 Jul 2026 16:45:14 +0800 Subject: [PATCH 4/4] fix: reorder KernelTable --- src/turbo/turbo.cc | 99 +++++++++++++++++++++++----------------------- 1 file changed, 50 insertions(+), 49 deletions(-) diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index c4394d0df..5402ad109 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -77,75 +77,76 @@ using RawBatchDistanceFn = void (*)(const void *const *, const void *, size_t, //! One row = one kernel family: all functions that must be used together //! for a given (metric, data type) combination on a given ISA. struct KernelSet { - MetricType metric; - DataType dtype; QuantizeType quantize; //!< QuantizeType served by this row - CpuArchType arch; //!< kScalar rows are the universal fallback + DataType dtype; + CpuArchType arch; //!< kScalar rows are the universal fallback + MetricType metric; RawDistanceFn dist; RawBatchDistanceFn batch; QueryPreprocessFunc preprocess; //!< non-null: batch needs preprocessing }; -// Dispatch registry. Row order encodes priority: SIMD rows must precede -// their scalar fallbacks. +// Dispatch registry, SIMD rows before their scalar +// fallbacks (row order encodes priority), then metric in enum order. constexpr KernelSet kKernelTable[] = { - // --- fp32 (scalar) --- - {MetricType::kCosine, DataType::kFp32, QuantizeType::kFp32, - CpuArchType::kScalar, scalar::cosine_fp32_distance, - scalar::cosine_fp32_batch_distance, nullptr}, - {MetricType::kSquaredEuclidean, DataType::kFp32, QuantizeType::kFp32, - CpuArchType::kScalar, scalar::squared_euclidean_fp32_distance, - scalar::squared_euclidean_fp32_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kFp32, QuantizeType::kFp32, - CpuArchType::kScalar, scalar::inner_product_fp32_distance, - scalar::inner_product_fp32_batch_distance, nullptr}, - - // --- fp16 (scalar) --- - {MetricType::kCosine, DataType::kFp16, QuantizeType::kFp16, - CpuArchType::kScalar, scalar::cosine_fp16_distance, - scalar::cosine_fp16_batch_distance, nullptr}, - {MetricType::kSquaredEuclidean, DataType::kFp16, QuantizeType::kFp16, - CpuArchType::kScalar, scalar::squared_euclidean_fp16_distance, - scalar::squared_euclidean_fp16_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kFp16, QuantizeType::kFp16, - CpuArchType::kScalar, scalar::inner_product_fp16_distance, - scalar::inner_product_fp16_batch_distance, nullptr}, - // --- record-quantized int8 (AVX512-VNNI, then scalar fallback) --- - {MetricType::kSquaredEuclidean, DataType::kInt8, QuantizeType::kRecord, - CpuArchType::kAVX512VNNI, avx512_vnni::squared_euclidean_int8_distance, + {QuantizeType::kRecord, DataType::kInt8, CpuArchType::kAVX512VNNI, + MetricType::kSquaredEuclidean, + avx512_vnni::squared_euclidean_int8_distance, avx512_vnni::squared_euclidean_int8_batch_distance, avx512_vnni::squared_euclidean_int8_query_preprocess}, - {MetricType::kCosine, DataType::kInt8, QuantizeType::kRecord, - CpuArchType::kAVX512VNNI, avx512_vnni::cosine_int8_distance, + {QuantizeType::kRecord, DataType::kInt8, CpuArchType::kAVX512VNNI, + MetricType::kCosine, avx512_vnni::cosine_int8_distance, avx512_vnni::cosine_int8_batch_distance, avx512_vnni::cosine_int8_query_preprocess}, - {MetricType::kSquaredEuclidean, DataType::kInt8, QuantizeType::kRecord, - CpuArchType::kScalar, scalar::squared_euclidean_int8_distance, + {QuantizeType::kRecord, DataType::kInt8, CpuArchType::kScalar, + MetricType::kSquaredEuclidean, scalar::squared_euclidean_int8_distance, scalar::squared_euclidean_int8_batch_distance, nullptr}, - {MetricType::kCosine, DataType::kInt8, QuantizeType::kRecord, - CpuArchType::kScalar, scalar::cosine_int8_distance, + {QuantizeType::kRecord, DataType::kInt8, CpuArchType::kScalar, + MetricType::kCosine, scalar::cosine_int8_distance, scalar::cosine_int8_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kInt8, QuantizeType::kRecord, - CpuArchType::kScalar, scalar::inner_product_int8_distance, + {QuantizeType::kRecord, DataType::kInt8, CpuArchType::kScalar, + MetricType::kInnerProduct, scalar::inner_product_int8_distance, scalar::inner_product_int8_batch_distance, nullptr}, - // --- uniform-quantized int8 (AVX512-VNNI only) --- - {MetricType::kSquaredEuclidean, DataType::kInt8, QuantizeType::kUniform, - CpuArchType::kAVX512VNNI, - avx512_vnni::uniform_squared_euclidean_int8_distance, - avx512_vnni::uniform_squared_euclidean_int8_batch_distance, nullptr}, - // --- record-quantized int4 (scalar) --- - {MetricType::kSquaredEuclidean, DataType::kInt4, QuantizeType::kRecord, - CpuArchType::kScalar, scalar::squared_euclidean_int4_distance, + {QuantizeType::kRecord, DataType::kInt4, CpuArchType::kScalar, + MetricType::kSquaredEuclidean, scalar::squared_euclidean_int4_distance, scalar::squared_euclidean_int4_batch_distance, nullptr}, - {MetricType::kCosine, DataType::kInt4, QuantizeType::kRecord, - CpuArchType::kScalar, scalar::cosine_int4_distance, + {QuantizeType::kRecord, DataType::kInt4, CpuArchType::kScalar, + MetricType::kCosine, scalar::cosine_int4_distance, scalar::cosine_int4_batch_distance, nullptr}, - {MetricType::kInnerProduct, DataType::kInt4, QuantizeType::kRecord, - CpuArchType::kScalar, scalar::inner_product_int4_distance, + {QuantizeType::kRecord, DataType::kInt4, CpuArchType::kScalar, + MetricType::kInnerProduct, scalar::inner_product_int4_distance, scalar::inner_product_int4_batch_distance, nullptr}, + + // --- uniform-quantized int8 (AVX512-VNNI only) --- + {QuantizeType::kUniform, DataType::kInt8, CpuArchType::kAVX512VNNI, + MetricType::kSquaredEuclidean, + avx512_vnni::uniform_squared_euclidean_int8_distance, + avx512_vnni::uniform_squared_euclidean_int8_batch_distance, nullptr}, + + // --- fp16 (scalar) --- + {QuantizeType::kFp16, DataType::kFp16, CpuArchType::kScalar, + MetricType::kSquaredEuclidean, scalar::squared_euclidean_fp16_distance, + scalar::squared_euclidean_fp16_batch_distance, nullptr}, + {QuantizeType::kFp16, DataType::kFp16, CpuArchType::kScalar, + MetricType::kCosine, scalar::cosine_fp16_distance, + scalar::cosine_fp16_batch_distance, nullptr}, + {QuantizeType::kFp16, DataType::kFp16, CpuArchType::kScalar, + MetricType::kInnerProduct, scalar::inner_product_fp16_distance, + scalar::inner_product_fp16_batch_distance, nullptr}, + + // --- fp32 (scalar) --- + {QuantizeType::kFp32, DataType::kFp32, CpuArchType::kScalar, + MetricType::kSquaredEuclidean, scalar::squared_euclidean_fp32_distance, + scalar::squared_euclidean_fp32_batch_distance, nullptr}, + {QuantizeType::kFp32, DataType::kFp32, CpuArchType::kScalar, + MetricType::kCosine, scalar::cosine_fp32_distance, + scalar::cosine_fp32_batch_distance, nullptr}, + {QuantizeType::kFp32, DataType::kFp32, CpuArchType::kScalar, + MetricType::kInnerProduct, scalar::inner_product_fp32_distance, + scalar::inner_product_fp32_batch_distance, nullptr}, }; // Returns the first (highest-priority) matching kernel row, or nullptr.