xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: enable icelake and l2 batch distance for int8 quantization. (#213)

This commit is contained in:
luoxiaojian 2026-03-13 17:26:22 +08:00 committed by GitHub
parent d1e0c7b88e
commit 79387c237f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 71 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
cmake/option.cmake
View File

@ -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()

6
src/core/metric/quantized_integer_metric.cc
View File

@ -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
} // namespace zvec

47
src/core/metric/quantized_integer_metric_batch.h
View File

@ -55,6 +55,11 @@ struct BaseDistanceBatchWithScoreUnquantized {
return CosineMinusInnerProductDistanceBatchWithScoreUnquantized<
ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim,
out);
} else if constexpr (std::is_same_v<DistanceType<ValueType, 1, 1>,
SquaredEuclidean<ValueType, 1, 1>>) {
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<int8_t, BatchSize,
float ms = m_tail[2];
float &result = results[i];
if (ImplType::GetQueryPreprocessFunc() != nullptr) {
int int_sum = reinterpret_cast<const int *>(m_tail)[3];
int int_sum = reinterpret_cast<const int *>(m_tail)[4];
result -= 128 * int_sum;
}
result = -(ma * qa * result + mb * qa * qs + qb * ma * ms +
@ -192,7 +197,7 @@ struct SquaredEuclideanDistanceBatchWithScoreUnquantized<int8_t, BatchSize,
PrefetchStep>;
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<int8_t, BatchSize,
const float sum = qa * qs;
const float sum2 = qa * qa * qs2;
for (int i = 0; i < num_vecs; ++i) {
for (size_t i = 0; i < num_vecs; ++i) {
const float *m_tail = reinterpret_cast<const float *>(
reinterpret_cast<const uint8_t *>(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<const int *>(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<int8_t, BatchSize,
}
static void QueryPreprocess(void *query, size_t dim) {
return ImplType::QueryPreprocess(query, dim - 16);
if (auto func = ImplType::GetQueryPreprocessFunc(); func != nullptr) {
return func(query, dim - 20);
}
}
};
@ -234,7 +245,7 @@ struct SquaredEuclideanDistanceBatchWithScoreUnquantized<int8_t, BatchSize,
template <size_t BatchSize, size_t PrefetchStep>
struct SquaredEuclideanDistanceBatchWithScoreUnquantized<uint8_t, BatchSize,
PrefetchStep> {
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<uint8_t, BatchSize,
const float sum = qa * qs;
const float sum2 = qa * qa * qs2;
for (int i = 0; i < num_vecs; ++i) {
for (size_t i = 0; i < num_vecs; ++i) {
const float *m_tail = reinterpret_cast<const float *>(
reinterpret_cast<const uint8_t *>(vecs[i]) +
original_dim_in_uint8_array);
@ -281,7 +292,7 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized<int8_t, BatchSize,
PrefetchStep>;
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<int8_t, BatchSize,
const float sum = qa * qs;
const float sum2 = qa * qa * qs2;
for (int i = 0; i < num_vecs; ++i) {
for (size_t i = 0; i < num_vecs; ++i) {
const float *m_tail = reinterpret_cast<const float *>(
reinterpret_cast<const int8_t *>(vecs[i]) + original_dim);
float ma = m_tail[0];
@ -310,7 +321,9 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized<int8_t, BatchSize,
}
static void QueryPreprocess(void *query, size_t dim) {
return ImplType::QueryPreprocess(query, dim - 16);
if (auto func = ImplType::GetQueryPreprocessFunc(); func != nullptr) {
return func(query, dim - 20);
}
}
};
@ -335,7 +348,7 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized<uint8_t, BatchSize,
const float sum = qa * qs;
const float sum2 = qa * qa * qs2;
for (int i = 0; i < num_vecs; ++i) {
for (size_t i = 0; i < num_vecs; ++i) {
const float *m_tail = reinterpret_cast<const float *>(
reinterpret_cast<const uint8_t *>(vecs[i]) +
original_dim_in_uint8_array);
@ -351,4 +364,4 @@ struct MipsSquaredEuclideanDistanceBatchWithScoreUnquantized<uint8_t, BatchSize,
}
};
} // namespace zvec::core
} // namespace zvec::core

10
src/core/metric/quantized_integer_metric_matrix.h
View File

@ -36,7 +36,7 @@ template <size_t M, size_t N>
struct SquaredEuclidean<int8_t, M, N> {
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<int8_t, M, N>::Compute(m, q, d, out);
for (size_t i = 0; i < N; ++i) {
@ -141,7 +141,7 @@ template <size_t M, size_t N>
struct MinusInnerProduct<int8_t, M, N> {
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, N>(m, q, origin_dim, out);
}
};
@ -168,7 +168,7 @@ template <size_t M, size_t N>
struct CosineMinusInnerProduct<int8_t, M, N> {
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, N>(m, q, origin_dim, out);
}
};
@ -195,7 +195,7 @@ template <size_t M, size_t N>
struct MipsSquaredEuclidean<int8_t, M, N> {
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<int8_t, M, N>::Compute(m, q, d, out);
for (size_t i = 0; i < N; ++i) {
@ -251,4 +251,4 @@ struct MipsSquaredEuclidean<uint8_t, M, N> {
}
};
} // namespace zvec::core
} // namespace zvec::core

6
src/core/quantizer/cosine_converter.cc
View File

@ -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
} // namespace zvec

2
src/core/quantizer/cosine_reformer.cc
View File

@ -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) {

5
src/core/quantizer/integer_quantizer_converter.cc
View File

@ -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

2
src/core/quantizer/integer_quantizer_reformer.cc
View File

@ -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

12
src/core/quantizer/record_quantizer.h
View File

@ -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<int32_t *>(extras + 4)[0] = int8_sum;
} else {
reinterpret_cast<int *>(extras)[3] = int8_sum;
if (is_euclidean) {
extras[3] = squared_sum;
} else {
reinterpret_cast<int *>(extras)[3] = int8_sum;
}
}
}
}
@ -128,4 +134,4 @@ class RecordQuantizer {
};
} // namespace core
} // namespace zvec
} // namespace zvec

8
tests/core/metric/quantized_integer_metric_test.cc
View File

@ -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<size_t>(1, 65))(gen) * 4;
size_t dimension = (std::uniform_int_distribution<size_t>(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<int8_t> matrix1(matrix_size);
std::vector<int8_t> 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<int8_t> query1(query_matrix_size);
std::vector<int8_t> 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<size_t>(1, 65))(gen) * 8;
size_t dimension = (std::uniform_int_distribution<size_t>(1, 65))(gen)*8;
auto holder = GetHolder(dimension, batch_size, dist);
IndexMeta meta(IndexMeta::DT_FP32, dimension);
meta.set_metric(metric_name, 0, Params());