From 79387c237fa5b65a747ea8a79b241b9bc7fc3056 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Fri, 13 Mar 2026 17:26:22 +0800 Subject: [PATCH] feat: enable icelake and l2 batch distance for int8 quantization. (#213) --- cmake/option.cmake | 12 +++-- src/core/metric/quantized_integer_metric.cc | 6 ++- .../metric/quantized_integer_metric_batch.h | 47 ++++++++++++------- .../metric/quantized_integer_metric_matrix.h | 10 ++-- src/core/quantizer/cosine_converter.cc | 6 +-- src/core/quantizer/cosine_reformer.cc | 2 +- .../quantizer/integer_quantizer_converter.cc | 5 +- .../quantizer/integer_quantizer_reformer.cc | 2 +- src/core/quantizer/record_quantizer.h | 12 +++-- .../metric/quantized_integer_metric_test.cc | 8 ++-- 10 files changed, 71 insertions(+), 39 deletions(-) diff --git a/cmake/option.cmake b/cmake/option.cmake index 71e4578..3c04242 100644 --- a/cmake/option.cmake +++ b/cmake/option.cmake @@ -9,6 +9,7 @@ option(ENABLE_HASWELL "Enable Intel Haswell CPU microarchitecture" OFF) option(ENABLE_BROADWELL "Enable Intel Broadwell CPU microarchitecture" OFF) option(ENABLE_SKYLAKE "Enable Intel Skylake CPU microarchitecture" OFF) option(ENABLE_SKYLAKE_AVX512 "Enable Intel Skylake Server CPU microarchitecture" OFF) +option(ENABLE_ICELAKE "Enable Intel Icelake CPU microarchitecture" OFF) option(ENABLE_SAPPHIRERAPIDS "Enable Intel Sapphire Rapids Server CPU microarchitecture" OFF) option(ENABLE_EMERALDRAPIDS "Enable Intel Emerald Rapids Server CPU microarchitecture" OFF) option(ENABLE_GRANITERAPIDS "Enable Intel Granite Rapids Server CPU microarchitecture" OFF) @@ -34,8 +35,8 @@ option(ENABLE_OPENMP "Enable OpenMP support" OFF) set(ARCH_OPTIONS ENABLE_NEHALEM ENABLE_SANDYBRIDGE ENABLE_HASWELL ENABLE_BROADWELL ENABLE_SKYLAKE - ENABLE_SKYLAKE_AVX512 ENABLE_SAPPHIRERAPIDS ENABLE_EMERALDRAPIDS ENABLE_GRANITERAPIDS - ENABLE_ZEN1 ENABLE_ZEN2 ENABLE_ZEN3 + ENABLE_SKYLAKE_AVX512 ENABLE_ICELAKE ENABLE_SAPPHIRERAPIDS ENABLE_EMERALDRAPIDS + ENABLE_GRANITERAPIDS ENABLE_ZEN1 ENABLE_ZEN2 ENABLE_ZEN3 ENABLE_ARMV8A ENABLE_ARMV8.1A ENABLE_ARMV8.2A ENABLE_ARMV8.3A ENABLE_ARMV8.4A ENABLE_ARMV8.5A ENABLE_ARMV8.6A ENABLE_NATIVE @@ -111,7 +112,8 @@ function(setup_compiler_march_for_x86 VAR_NAME_SSE VAR_NAME_AVX2 VAR_NAME_AVX512 #avx512 set(_x86_flags - "graniterapids" "emeraldrapids" "sapphirerapids" "skylake-avx512" + "graniterapids" "emeraldrapids" "sapphirerapids" + "icelake-server" "skylake-avx512" ) foreach(_arch IN LISTS _x86_flags) check_c_compiler_flag("-march=${_arch}" _COMP_SUPP_${_arch}) @@ -170,6 +172,10 @@ if(NOT AUTO_DETECT_ARCH) add_arch_flag("-march=sapphirerapids" SAPPHIRERAPIDS ENABLE_SAPPHIRERAPIDS) endif() + if(ENABLE_ICELAKE) + add_arch_flag("-march=icelake-server" ICELAKE ENABLE_ICELAKE) + endif() + if(ENABLE_SKYLAKE_AVX512) add_arch_flag("-march=skylake-avx512" SKYLAKE_AVX512 ENABLE_SKYLAKE_AVX512) endif() diff --git a/src/core/metric/quantized_integer_metric.cc b/src/core/metric/quantized_integer_metric.cc index 56e9563..2b4e757 100644 --- a/src/core/metric/quantized_integer_metric.cc +++ b/src/core/metric/quantized_integer_metric.cc @@ -266,6 +266,10 @@ class QuantizedIntegerMetric : public IndexMetric { meta_.data_type() == IndexMeta::DataType::DT_INT8) { return CosineMinusInnerProductDistanceBatchWithScoreUnquantized< int8_t, 1, 1>::GetQueryPreprocessFunc(); + } else if (origin_metric_type_ == MetricType::kSquaredEuclidean && + meta_.data_type() == IndexMeta::DataType::DT_INT8) { + return SquaredEuclideanDistanceBatchWithScoreUnquantized< + int8_t, 1, 1>::GetQueryPreprocessFunc(); } return nullptr; @@ -320,4 +324,4 @@ class QuantizedIntegerMetric : public IndexMetric { INDEX_FACTORY_REGISTER_METRIC_ALIAS(QuantizedInteger, QuantizedIntegerMetric); } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/metric/quantized_integer_metric_batch.h b/src/core/metric/quantized_integer_metric_batch.h index e9e63ce..33bbfa9 100644 --- a/src/core/metric/quantized_integer_metric_batch.h +++ b/src/core/metric/quantized_integer_metric_batch.h @@ -55,6 +55,11 @@ struct BaseDistanceBatchWithScoreUnquantized { return CosineMinusInnerProductDistanceBatchWithScoreUnquantized< ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim, out); + } else if constexpr (std::is_same_v, + SquaredEuclidean>) { + return SquaredEuclideanDistanceBatchWithScoreUnquantized< + ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim, + out); } _ComputeBatch(m, q, num, dim, out); @@ -75,7 +80,7 @@ struct CosineMinusInnerProductDistanceBatchWithScoreUnquantized< static inline void ComputeBatch(const int8_t **vecs, const int8_t *query, size_t num_vecs, size_t dim, float *results) { - size_t original_dim = dim - 20; + size_t original_dim = dim - 24; ImplType::ComputeBatch(vecs, query, num_vecs, original_dim, results); } @@ -87,7 +92,7 @@ struct CosineMinusInnerProductDistanceBatchWithScoreUnquantized< static void QueryPreprocess(void *query, size_t dim) { if (auto func = ImplType::GetQueryPreprocessFunc(); func != nullptr) { - return func(query, dim - 20); + return func(query, dim - 24); } } }; @@ -134,7 +139,7 @@ struct MinusInnerProductDistanceBatchWithScoreUnquantized(m_tail)[3]; + int int_sum = reinterpret_cast(m_tail)[4]; result -= 128 * int_sum; } result = -(ma * qa * result + mb * qa * qs + qb * ma * ms + @@ -192,7 +197,7 @@ struct SquaredEuclideanDistanceBatchWithScoreUnquantized; static void ComputeBatch(const int8_t **vecs, const int8_t *query, size_t num_vecs, size_t dim, float *results) { - const size_t original_dim = dim - 16; + const size_t original_dim = dim - 20; ailego::DistanceBatch::InnerProductDistanceBatch< int8_t, BatchSize, PrefetchStep>::ComputeBatch(vecs, query, num_vecs, original_dim, results); @@ -206,17 +211,21 @@ struct SquaredEuclideanDistanceBatchWithScoreUnquantized( reinterpret_cast(vecs[i]) + original_dim); float ma = m_tail[0]; float mb = m_tail[1]; float ms = m_tail[2]; float ms2 = m_tail[3]; - *results = ma * ma * ms2 + sum2 - 2 * ma * qa * *results + - (mb - qb) * (mb - qb) * original_dim + - 2 * (mb - qb) * (ms * ma - sum); - ++results; + float &result = results[i]; + if (ImplType::GetQueryPreprocessFunc() != nullptr) { + int int8_sum = reinterpret_cast(m_tail)[4]; + result -= 128 * int8_sum; + } + result = ma * ma * ms2 + sum2 - 2 * ma * qa * result + + (mb - qb) * (mb - qb) * original_dim + + 2 * (mb - qb) * (ms * ma - sum); } } @@ -226,7 +235,9 @@ struct SquaredEuclideanDistanceBatchWithScoreUnquantized struct SquaredEuclideanDistanceBatchWithScoreUnquantized { - static void ComputeBatch(const int8_t **vecs, const int8_t *query, + static void ComputeBatch(const uint8_t **vecs, const uint8_t *query, size_t num_vecs, size_t dim, float *results) { const size_t original_dim = dim - 32; const size_t original_dim_in_uint8_array = original_dim >> 1; @@ -251,7 +262,7 @@ struct SquaredEuclideanDistanceBatchWithScoreUnquantized( reinterpret_cast(vecs[i]) + original_dim_in_uint8_array); @@ -281,7 +292,7 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized; static void ComputeBatch(const int8_t **vecs, const int8_t *query, size_t num_vecs, size_t dim, float *results) { - const size_t original_dim = dim - 16; + const size_t original_dim = dim - 20; ailego::DistanceBatch::InnerProductDistanceBatch< int8_t, BatchSize, PrefetchStep>::ComputeBatch(vecs, query, num_vecs, original_dim, results); @@ -295,7 +306,7 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized( reinterpret_cast(vecs[i]) + original_dim); float ma = m_tail[0]; @@ -310,7 +321,9 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized( reinterpret_cast(vecs[i]) + original_dim_in_uint8_array); @@ -351,4 +364,4 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized struct SquaredEuclidean { static void Compute(const int8_t *m, const int8_t *q, size_t dim, float *out) { - const size_t d = dim - 16; + const size_t d = dim - 20; ailego::InnerProductMatrix::Compute(m, q, d, out); for (size_t i = 0; i < N; ++i) { @@ -141,7 +141,7 @@ template struct MinusInnerProduct { static void Compute(const int8_t *m, const int8_t *q, size_t dim, float *out) { - const size_t origin_dim = dim - 16; + const size_t origin_dim = dim - 20; MinusInnerProductImplInt8(m, q, origin_dim, out); } }; @@ -168,7 +168,7 @@ template struct CosineMinusInnerProduct { static void Compute(const int8_t *m, const int8_t *q, size_t dim, float *out) { - const size_t origin_dim = dim - 20; + const size_t origin_dim = dim - 24; MinusInnerProductImplInt8(m, q, origin_dim, out); } }; @@ -195,7 +195,7 @@ template struct MipsSquaredEuclidean { static void Compute(const int8_t *m, const int8_t *q, size_t dim, float *out) { - const size_t d = dim - 16; + const size_t d = dim - 20; ailego::InnerProductMatrix::Compute(m, q, d, out); for (size_t i = 0; i < N; ++i) { @@ -251,4 +251,4 @@ struct MipsSquaredEuclidean { } }; -} // namespace zvec::core \ No newline at end of file +} // namespace zvec::core diff --git a/src/core/quantizer/cosine_converter.cc b/src/core/quantizer/cosine_converter.cc index dda76b0..dd5cbbd 100644 --- a/src/core/quantizer/cosine_converter.cc +++ b/src/core/quantizer/cosine_converter.cc @@ -206,7 +206,7 @@ class CosineConverterHolder : public IndexHolder { if (type == IndexMeta::DataType::DT_INT4) return 40; // 5 * sizeof(float) / sizeof(FT_INT4) else if (type == IndexMeta::DataType::DT_INT8) - return 20; // 5 * sizeof(float) / sizeof(FT_INT8) + return 24; // (5 * sizeof(float) + sizeof(int)) / sizeof(FT_INT8) else if (type == IndexMeta::DataType::DT_FP16) return 2; // 2* sizeof(float) / sizeof(FT_FP16) else if (type == IndexMeta::DataType::DT_FP32) { @@ -362,7 +362,7 @@ class CosineConverter : public IndexConverter { if (type == IndexMeta::DataType::DT_INT4) return 40; // 5 * sizeof(float) / sizeof(FT_INT4) else if (type == IndexMeta::DataType::DT_INT8) - return 20; // 5 * sizeof(float) / sizeof(FT_INT8) + return 24; // (5 * sizeof(float) + sizeof(int)) / sizeof(FT_INT8) else if (type == IndexMeta::DataType::DT_FP16) return 2; // sizeof(float) / sizeof(FT_FP16) else if (type == IndexMeta::DataType::DT_FP32) { @@ -402,4 +402,4 @@ INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(CosineHalfFloatConverter, IndexMeta::DataType::DT_FP16); } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/quantizer/cosine_reformer.cc b/src/core/quantizer/cosine_reformer.cc index 5823728..d6080b8 100644 --- a/src/core/quantizer/cosine_reformer.cc +++ b/src/core/quantizer/cosine_reformer.cc @@ -249,7 +249,7 @@ class CosineReformer : public IndexReformer { if (type == IndexMeta::DataType::DT_INT4) return 40; // 5 * sizeof(float) / sizeof(FT_INT4) else if (type == IndexMeta::DataType::DT_INT8) - return 20; // 5 * sizeof(float) / sizeof(FT_INT8) + return 24; // (5 * sizeof(float) + sizeof(int)) / sizeof(FT_INT8) else if (type == IndexMeta::DataType::DT_FP16) return 2; // sizeof(float) / sizeof(FT_FP16) else if (type == IndexMeta::DataType::DT_FP32) { diff --git a/src/core/quantizer/integer_quantizer_converter.cc b/src/core/quantizer/integer_quantizer_converter.cc index 91757a5..1cd8884 100644 --- a/src/core/quantizer/integer_quantizer_converter.cc +++ b/src/core/quantizer/integer_quantizer_converter.cc @@ -581,7 +581,10 @@ class IntegerStreamingConverter : public IndexConverter { static size_t ExtraDimension(IndexMeta::DataType type) { // The extra quantized params storage size to save for each vector constexpr size_t kExtraSize = 4 * sizeof(float); - return type == IndexMeta::DataType::DT_INT8 ? kExtraSize : kExtraSize * 2; + constexpr size_t kAdditionalInt32 = sizeof(int32_t); + return type == IndexMeta::DataType::DT_INT8 + ? (kExtraSize + kAdditionalInt32) + : (kExtraSize * 2); } //! Members diff --git a/src/core/quantizer/integer_quantizer_reformer.cc b/src/core/quantizer/integer_quantizer_reformer.cc index 9c74103..4228d0f 100644 --- a/src/core/quantizer/integer_quantizer_reformer.cc +++ b/src/core/quantizer/integer_quantizer_reformer.cc @@ -279,7 +279,7 @@ class IntegerStreamingReformer : public IndexReformer { //! Constructor IntegerStreamingReformer(IndexMeta::DataType dst_type) : data_type_(dst_type), - extra_dimension_(data_type_ == IndexMeta::DataType::DT_INT8 ? 16 : 32) { + extra_dimension_(data_type_ == IndexMeta::DataType::DT_INT8 ? 20 : 32) { } //! Initialize Reformer diff --git a/src/core/quantizer/record_quantizer.h b/src/core/quantizer/record_quantizer.h index 06744f6..b1095a2 100644 --- a/src/core/quantizer/record_quantizer.h +++ b/src/core/quantizer/record_quantizer.h @@ -74,10 +74,16 @@ class RecordQuantizer { extras[0] = 1.0f / scale; extras[1] = -bias / scale; extras[2] = sum; - if (is_euclidean) { + + if (type == IndexMeta::DataType::DT_INT8) { extras[3] = squared_sum; + reinterpret_cast(extras + 4)[0] = int8_sum; } else { - reinterpret_cast(extras)[3] = int8_sum; + if (is_euclidean) { + extras[3] = squared_sum; + } else { + reinterpret_cast(extras)[3] = int8_sum; + } } } } @@ -128,4 +134,4 @@ class RecordQuantizer { }; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/tests/core/metric/quantized_integer_metric_test.cc b/tests/core/metric/quantized_integer_metric_test.cc index 30e8c25..501d8c7 100644 --- a/tests/core/metric/quantized_integer_metric_test.cc +++ b/tests/core/metric/quantized_integer_metric_test.cc @@ -251,7 +251,7 @@ void TestDistanceMatrixInt8(const std::string &metric_name) { const size_t batch_size = M; const size_t query_size = N; - size_t dimension = (std::uniform_int_distribution(1, 65))(gen) * 4; + size_t dimension = (std::uniform_int_distribution(1, 65))(gen)*4; auto holder = GetHolder(dimension, batch_size, dist); IndexMeta meta(IndexMeta::DT_FP32, dimension); meta.set_metric(metric_name, 0, Params()); @@ -261,7 +261,7 @@ void TestDistanceMatrixInt8(const std::string &metric_name) { ASSERT_EQ(0u, IndexConverter::TrainAndTransform(converter, holder)); auto holder2 = converter->result(); auto &meta2 = converter->meta(); - ASSERT_EQ(dimension + 16, holder2->dimension()); + ASSERT_EQ(dimension + 20, holder2->dimension()); size_t matrix_size = batch_size * holder2->dimension(); std::vector matrix1(matrix_size); std::vector matrix2(matrix_size); @@ -277,7 +277,7 @@ void TestDistanceMatrixInt8(const std::string &metric_name) { auto query_holder = GetHolder(dimension, query_size, dist); ASSERT_EQ(0u, IndexConverter::TrainAndTransform(converter, query_holder)); auto query_holder2 = converter->result(); - ASSERT_EQ(dimension + 16, query_holder2->dimension()); + ASSERT_EQ(dimension + 20, query_holder2->dimension()); size_t query_matrix_size = query_size * query_holder2->dimension(); std::vector query1(query_matrix_size); std::vector query2(query_matrix_size); @@ -453,7 +453,7 @@ void TestDistanceMatrixInt4(const std::string &metric_name) { const size_t batch_size = M; const size_t query_size = N; - size_t dimension = (std::uniform_int_distribution(1, 65))(gen) * 8; + size_t dimension = (std::uniform_int_distribution(1, 65))(gen)*8; auto holder = GetHolder(dimension, batch_size, dist); IndexMeta meta(IndexMeta::DT_FP32, dimension); meta.set_metric(metric_name, 0, Params());