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

feat: add euclidean one2many (#188) 

* add euclidean one2many implementation
This commit is contained in:
rayx 2026-03-23 20:20:48 +08:00 committed by GitHub
parent 22cb3cc4b5
commit 31aaa2febb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 983 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
src/ailego/math/matrix_utility.i
View File

@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <zvec/ailego/internal/platform.h>
namespace zvec {
@ -46,7 +48,7 @@ static inline float HorizontalAdd_FP32_V128(__m128 v) {
return _mm_cvtss_f32(x4);
#endif
}
#endif // __SSE__
#endif // __SSE__
#if defined(__SSE2__)
static inline int32_t HorizontalAdd_INT32_V128(__m128i v) {
@ -71,7 +73,7 @@ static inline int64_t HorizontalAdd_INT64_V128(__m128i v) {
_mm_add_epi64(_mm_shuffle_epi32(v, _MM_SHUFFLE(0, 0, 3, 2)), v));
#endif
}
#endif // __SSE2__
#endif // __SSE2__
#if defined(__SSSE3__)
static const __m128i POPCNT_LOOKUP_SSE =
@ -86,7 +88,7 @@ static inline __m128i VerticalPopCount_INT8_V128(__m128i v) {
__m128i hi = _mm_shuffle_epi8(POPCNT_LOOKUP_SSE,
_mm_and_si128(_mm_srli_epi32(v, 4), low_mask));
return _mm_add_epi8(lo, hi);
#endif // __AVX512VL__ && __AVX512BITALG__
#endif // __AVX512VL__ && __AVX512BITALG__
}
static inline __m128i VerticalPopCount_INT16_V128(__m128i v) {
@ -96,7 +98,7 @@ static inline __m128i VerticalPopCount_INT16_V128(__m128i v) {
__m128i total = VerticalPopCount_INT8_V128(v);
return _mm_add_epi16(_mm_srli_epi16(total, 8),
_mm_and_si128(total, _mm_set1_epi16(0xff)));
#endif // __AVX512VL__ && __AVX512BITALG__
#endif // __AVX512VL__ && __AVX512BITALG__
}
static inline __m128i VerticalPopCount_INT32_V128(__m128i v) {
@ -107,7 +109,7 @@ static inline __m128i VerticalPopCount_INT32_V128(__m128i v) {
_mm_madd_epi16(VerticalPopCount_INT8_V128(v), _mm_set1_epi16(1));
return _mm_add_epi32(_mm_srli_epi32(total, 8),
_mm_and_si128(total, _mm_set1_epi32(0xff)));
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
}
static inline __m128i VerticalPopCount_INT64_V128(__m128i v) {
@ -115,9 +117,9 @@ static inline __m128i VerticalPopCount_INT64_V128(__m128i v) {
return _mm_popcnt_epi64(v);
#else
return _mm_sad_epu8(VerticalPopCount_INT8_V128(v), _mm_setzero_si128());
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
}
#endif // __SSSE3__
#endif // __SSSE3__
#if defined(__SSE4_1__)
static inline int16_t HorizontalMax_UINT8_V128(__m128i v) {
@ -127,7 +129,7 @@ static inline int16_t HorizontalMax_UINT8_V128(__m128i v) {
v = _mm_max_epu8(v, _mm_srli_epi16(v, 8));
return static_cast<uint8_t>(_mm_cvtsi128_si32(v));
}
#endif // __SSE4_1__
#endif // __SSE4_1__
#if defined(__AVX__)
static inline float HorizontalMax_FP32_V256(__m256 v) {
@ -147,7 +149,7 @@ static inline float HorizontalAdd_FP32_V256(__m256 v) {
__m128 x4 = _mm_add_ss(_mm256_castps256_ps128(x2), x3);
return _mm_cvtss_f32(x4);
}
#endif // __AVX__
#endif // __AVX__
#if defined(__AVX2__)
#define POPCNT_MASK1_INT8_AVX _mm256_set1_epi8(0x0f)
@ -167,7 +169,7 @@ static inline __m256i VerticalPopCount_INT8_V256(__m256i v) {
POPCNT_LOOKUP_AVX,
_mm256_and_si256(_mm256_srli_epi32(v, 4), POPCNT_MASK1_INT8_AVX));
return _mm256_add_epi8(lo, hi);
#endif // __AVX512VL__ && __AVX512BITALG__
#endif // __AVX512VL__ && __AVX512BITALG__
}
static inline __m256i VerticalPopCount_INT16_V256(__m256i v) {
@ -177,7 +179,7 @@ static inline __m256i VerticalPopCount_INT16_V256(__m256i v) {
__m256i total = VerticalPopCount_INT8_V256(v);
return _mm256_add_epi16(_mm256_srli_epi16(total, 8),
_mm256_and_si256(total, POPCNT_MASK2_INT16_AVX));
#endif // __AVX512VL__ && __AVX512BITALG__
#endif // __AVX512VL__ && __AVX512BITALG__
}
static inline __m256i VerticalPopCount_INT32_V256(__m256i v) {
@ -188,7 +190,7 @@ static inline __m256i VerticalPopCount_INT32_V256(__m256i v) {
_mm256_madd_epi16(VerticalPopCount_INT8_V256(v), POPCNT_MASK1_INT16_AVX);
return _mm256_add_epi32(_mm256_srli_epi32(total, 8),
_mm256_and_si256(total, POPCNT_MASK1_INT32_AVX));
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
}
static inline __m256i VerticalPopCount_INT64_V256(__m256i v) {
@ -196,7 +198,7 @@ static inline __m256i VerticalPopCount_INT64_V256(__m256i v) {
return _mm256_popcnt_epi64(v);
#else
return _mm256_sad_epu8(VerticalPopCount_INT8_V256(v), POPCNT_ZERO_AVX);
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
#endif // __AVX512VL__ && __AVX512VPOPCNTDQ__
}
static inline int16_t HorizontalMax_UINT8_V256(__m256i v) {
@ -224,7 +226,7 @@ static inline int64_t HorizontalAdd_INT64_V256(__m256i v) {
__m128i x4 = _mm_add_epi64(_mm256_extractf128_si256(x2, 0), x3);
return _mm_cvtsi128_si64(x4);
}
#endif // __AVX2__
#endif // __AVX2__
#if defined(__AVX512F__)
static inline float HorizontalMax_FP32_V512(__m512 v) {
@ -240,7 +242,7 @@ static inline float HorizontalAdd_FP32_V512(__m512 v) {
_mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(v), 1));
return HorizontalAdd_FP32_V256(_mm256_add_ps(low, high));
}
#endif // __AVX512F__
#endif // __AVX512F__
#if defined(__AVX512FP16__)
static inline float HorizontalMax_FP16_V512(__m512h v) {
@ -257,7 +259,7 @@ static inline float HorizontalAdd_FP16_V512(__m512h v) {
return HorizontalAdd_FP32_V512(_mm512_add_ps(low, high));
}
#endif // __AVX512FP16__
#endif // __AVX512FP16__
} // namespace ailego
} // namespace zvec

16
src/ailego/math_batch/distance_batch.h
View File

@ -17,6 +17,7 @@
#include <zvec/ailego/math_batch/utils.h>
#include "ailego/math/distance_matrix.h"
#include "cosine_distance_batch.h"
#include "euclidean_distance_batch.h"
#include "inner_product_distance_batch.h"
namespace zvec::ailego {
@ -42,6 +43,21 @@ struct BaseDistance {
out);
}
if constexpr (std::is_same_v<DistanceType<ValueType, 1, 1>,
EuclideanDistanceMatrix<ValueType, 1, 1>>) {
return DistanceBatch::EuclideanDistanceBatch<
ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim,
out);
}
if constexpr (std::is_same_v<
DistanceType<ValueType, 1, 1>,
SquaredEuclideanDistanceMatrix<ValueType, 1, 1>>) {
return DistanceBatch::SquaredEuclideanDistanceBatch<
ValueType, BatchSize, PrefetchStep>::ComputeBatch(m, q, num, dim,
out);
}
_ComputeBatch(m, q, num, dim, out);
}
};

30
src/ailego/math_batch/distance_batch_math.h Normal file
View File

@ -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
#if defined(__AVX2__)
inline float sum4(__m128 v) {
v = _mm_add_ps(v, _mm_castsi128_ps(_mm_srli_si128(_mm_castps_si128(v), 8)));
return v[0] + v[1];
}
inline __m128 sum_top_bottom_avx(__m256 v) {
const __m128 high = _mm256_extractf128_ps(v, 1);
const __m128 low = _mm256_castps256_ps128(v);
return _mm_add_ps(high, low);
}
#endif

142
src/ailego/math_batch/euclidean_distance_batch.h Normal file
View File

@ -0,0 +1,142 @@
// 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 <vector>
#include <ailego/math/euclidean_distance_matrix.h>
#include <ailego/utility/math_helper.h>
#include <zvec/ailego/internal/platform.h>
#include <zvec/ailego/math_batch/utils.h>
#include <zvec/ailego/utility/type_helper.h>
namespace zvec::ailego::DistanceBatch {
// SquaredEuclideanDistanceBatch
template <typename T, size_t BatchSize, size_t PrefetchStep, typename = void>
struct SquaredEuclideanDistanceBatch;
template <typename ValueType, size_t BatchSize>
static void compute_one_to_many_squared_euclidean_fallback(
const ValueType *query, const ValueType **ptrs,
std::array<const ValueType *, BatchSize> &prefetch_ptrs, size_t dim,
float *sums) {
for (size_t j = 0; j < BatchSize; ++j) {
sums[j] = 0.0;
SquaredEuclideanDistanceMatrix<ValueType, 1, 1>::Compute(ptrs[j], query,
dim, sums + j);
ailego_prefetch(&prefetch_ptrs[j]);
}
}
// Function template partial specialization is not allowed,
// therefore the wrapper struct is required.
template <typename T, size_t BatchSize>
struct SquaredEuclideanDistanceBatchImpl {
using ValueType = typename std::remove_cv<T>::type;
static void compute_one_to_many(
const ValueType *query, const ValueType **ptrs,
std::array<const ValueType *, BatchSize> &prefetch_ptrs, size_t dim,
float *sums) {
return compute_one_to_many_squared_euclidean_fallback(
query, ptrs, prefetch_ptrs, dim, sums);
}
};
template <typename T, size_t BatchSize, size_t PrefetchStep, typename>
struct SquaredEuclideanDistanceBatch {
using ValueType = typename std::remove_cv<T>::type;
static inline void ComputeBatch(const ValueType **vecs,
const ValueType *query, size_t num_vecs,
size_t dim, float *results) {
size_t i = 0;
for (; i + BatchSize <= num_vecs; i += BatchSize) {
std::array<const ValueType *, BatchSize> prefetch_ptrs;
for (size_t j = 0; j < BatchSize; ++j) {
if (i + j + BatchSize * PrefetchStep < num_vecs) {
prefetch_ptrs[j] = vecs[i + j + BatchSize * PrefetchStep];
} else {
prefetch_ptrs[j] = nullptr;
}
}
SquaredEuclideanDistanceBatchImpl<
ValueType, BatchSize>::compute_one_to_many(query, &vecs[i],
prefetch_ptrs, dim,
&results[i]);
}
for (; i < num_vecs; ++i) { // TODO: unroll by 1, 2, 4, 8, etc.
std::array<const ValueType *, 1> prefetch_ptrs{nullptr};
SquaredEuclideanDistanceBatchImpl<ValueType, 1>::compute_one_to_many(
query, &vecs[i], prefetch_ptrs, dim, &results[i]);
}
}
};
// EuclideanDistanceBatch
template <typename T, size_t BatchSize, size_t PrefetchStep, typename = void>
struct EuclideanDistanceBatch;
template <typename T, size_t BatchSize, size_t PrefetchStep, typename>
struct EuclideanDistanceBatch {
using ValueType = typename std::remove_cv<T>::type;
static inline void ComputeBatch(const ValueType **vecs,
const ValueType *query, size_t num_vecs,
size_t dim, float *results) {
SquaredEuclideanDistanceBatch<T, BatchSize, PrefetchStep>::ComputeBatch(
vecs, query, num_vecs, dim, results);
for (size_t i = 0; i < num_vecs; ++i) {
results[i] = std::sqrt(results[i]);
}
}
};
template <>
struct SquaredEuclideanDistanceBatchImpl<ailego::Float16, 1> {
using ValueType = ailego::Float16;
static void compute_one_to_many(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs, size_t dim,
float *sums);
};
template <>
struct SquaredEuclideanDistanceBatchImpl<float, 1> {
using ValueType = float;
static void compute_one_to_many(const float *query, const float **ptrs,
std::array<const float *, 1> &prefetch_ptrs,
size_t dim, float *sums);
};
template <>
struct SquaredEuclideanDistanceBatchImpl<ailego::Float16, 12> {
using ValueType = ailego::Float16;
static void compute_one_to_many(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs, size_t dim,
float *sums);
};
template <>
struct SquaredEuclideanDistanceBatchImpl<float, 12> {
using ValueType = float;
static void compute_one_to_many(const float *query, const float **ptrs,
std::array<const float *, 12> &prefetch_ptrs,
size_t dim, float *sums);
};
} // namespace zvec::ailego::DistanceBatch

191
src/ailego/math_batch/euclidean_distance_batch_dispatch.cc Normal file
View File

@ -0,0 +1,191 @@
// 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 <ailego/internal/cpu_features.h>
#include <ailego/utility/math_helper.h>
#include <zvec/ailego/internal/platform.h>
#include <zvec/ailego/math_batch/utils.h>
#include <zvec/ailego/utility/type_helper.h>
#include "euclidean_distance_batch.h"
namespace zvec::ailego::DistanceBatch {
#if defined(__AVX512FP16__)
void compute_one_to_many_squared_euclidean_avx512fp16_fp16_1(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs,
size_t dimensionality, float *results);
void compute_one_to_many_squared_euclidean_avx512fp16_fp16_12(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs,
size_t dimensionality, float *results);
#endif //__AVX512FP16__
#if defined(__AVX512F__)
void compute_one_to_many_squared_euclidean_avx512f_fp16_1(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs,
size_t dimensionality, float *results);
void compute_one_to_many_squared_euclidean_avx512f_fp16_12(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs,
size_t dimensionality, float *results);
void compute_one_to_many_squared_euclidean_avx512f_fp32_1(
const float *query, const float **ptrs,
std::array<const float *, 1> &prefetch_ptrs, size_t dimensionality,
float *results);
void compute_one_to_many_squared_euclidean_avx512f_fp32_12(
const float *query, const float **ptrs,
std::array<const float *, 12> &prefetch_ptrs, size_t dimensionality,
float *results);
#endif //__AVX512F__
#if defined(__AVX2__)
void compute_one_to_many_squared_euclidean_avx2_fp32_1(
const float *query, const float **ptrs,
std::array<const float *, 1> &prefetch_ptrs, size_t dimensionality,
float *results);
void compute_one_to_many_squared_euclidean_avx2_fp16_1(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs,
size_t dimensionality, float *results);
// void compute_one_to_many_squared_euclidean_avx2_int8_1(
// const int8_t *query, const int8_t **ptrs,
// std::array<const int8_t *, 1> &prefetch_ptrs, size_t dimensionality,
// float *results);
void compute_one_to_many_squared_euclidean_avx2_fp32_12(
const float *query, const float **ptrs,
std::array<const float *, 12> &prefetch_ptrs, size_t dimensionality,
float *results);
void compute_one_to_many_squared_euclidean_avx2_fp16_12(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs,
size_t dimensionality, float *results);
// void compute_one_to_many_squared_euclidean_avx2_int8_12(
// const int8_t *query, const int8_t **ptrs,
// std::array<const int8_t *, 12> &prefetch_ptrs, size_t dimensionality,
// float *results);
#endif
void SquaredEuclideanDistanceBatchImpl<float, 1>::compute_one_to_many(
const ValueType *query, const ValueType **ptrs,
std::array<const ValueType *, 1> &prefetch_ptrs, size_t dim, float *sums) {
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
return compute_one_to_many_squared_euclidean_avx2_fp32_1(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
return compute_one_to_many_squared_euclidean_fallback(
query, ptrs, prefetch_ptrs, dim, sums);
}
void SquaredEuclideanDistanceBatchImpl<ailego::Float16, 1>::compute_one_to_many(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs, size_t dim,
float *sums) {
#if defined(__AVX512FP16__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {
return compute_one_to_many_squared_euclidean_avx512fp16_fp16_1(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
return compute_one_to_many_squared_euclidean_avx512f_fp16_1(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
return compute_one_to_many_squared_euclidean_avx2_fp16_1(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
return compute_one_to_many_squared_euclidean_fallback(
query, ptrs, prefetch_ptrs, dim, sums);
}
void SquaredEuclideanDistanceBatchImpl<float, 12>::compute_one_to_many(
const float *query, const float **ptrs,
std::array<const float *, 12> &prefetch_ptrs, size_t dim, float *sums) {
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
return compute_one_to_many_squared_euclidean_avx512f_fp32_12(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
return compute_one_to_many_squared_euclidean_avx2_fp32_12(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
return compute_one_to_many_squared_euclidean_fallback(
query, ptrs, prefetch_ptrs, dim, sums);
}
void SquaredEuclideanDistanceBatchImpl<ailego::Float16, 12>::
compute_one_to_many(const ailego::Float16 *query,
const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs,
size_t dim, float *sums) {
#if defined(__AVX512FP16__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_FP16) {
return compute_one_to_many_squared_euclidean_avx512fp16_fp16_12(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
return compute_one_to_many_squared_euclidean_avx512f_fp16_12(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
return compute_one_to_many_squared_euclidean_avx2_fp16_12(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
return compute_one_to_many_squared_euclidean_fallback(
query, ptrs, prefetch_ptrs, dim, sums);
}
// void SquaredEuclideanDistanceBatchImpl<int8_t, 12>::compute_one_to_many(
// const int8_t *query, const int8_t **ptrs,
// std::array<const int8_t *, 12> &prefetch_ptrs, size_t dim, float *sums) {
// #if defined(__AVX2__)
// if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
// return compute_one_to_many_squared_euclidean_avx2_int8_12(
// query, ptrs, prefetch_ptrs, dim, sums);
// }
// #endif
// return compute_one_to_many_squared_euclidean_fallback(query, ptrs,
// prefetch_ptrs,
// dim, sums);
// }
} // namespace zvec::ailego::DistanceBatch

114
src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx2.cc Normal file
View File

@ -0,0 +1,114 @@
// 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 <ailego/math/matrix_utility.i>
#include <ailego/utility/math_helper.h>
#include <zvec/ailego/internal/platform.h>
#include <zvec/ailego/utility/type_helper.h>
namespace zvec::ailego::DistanceBatch {
#if defined(__AVX2__)
template <typename ValueType, size_t dp_batch>
static std::enable_if_t<std::is_same_v<ValueType, ailego::Float16>, void>
compute_one_to_many_squared_euclidean_avx2_fp16(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, dp_batch> &prefetch_ptrs,
size_t dimensionality, float *results) {
__m256 accs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
accs[i] = _mm256_setzero_ps();
}
size_t dim = 0;
for (; dim + 16 <= dimensionality; dim += 16) {
__m256i q =
_mm256_loadu_si256(reinterpret_cast<const __m256i *>(query + dim));
__m256 q1 = _mm256_cvtph_ps(_mm256_castsi256_si128(q));
__m256 q2 = _mm256_cvtph_ps(_mm256_extractf128_si256(q, 1));
__m256 data_regs_1[dp_batch];
__m256 data_regs_2[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
__m256i m =
_mm256_loadu_si256(reinterpret_cast<const __m256i *>(ptrs[i] + dim));
data_regs_1[i] = _mm256_cvtph_ps(_mm256_castsi256_si128(m));
data_regs_2[i] = _mm256_cvtph_ps(_mm256_extractf128_si256(m, 1));
}
if (prefetch_ptrs[0]) {
for (size_t i = 0; i < dp_batch; ++i) {
ailego_prefetch(prefetch_ptrs[i] + dim);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
__m256 diff1 = _mm256_sub_ps(q1, data_regs_1[i]);
accs[i] = _mm256_fmadd_ps(diff1, diff1, accs[i]);
__m256 diff2 = _mm256_sub_ps(q2, data_regs_2[i]);
accs[i] = _mm256_fmadd_ps(diff2, diff2, accs[i]);
}
}
if (dim + 8 <= dimensionality) {
__m256 q = _mm256_cvtph_ps(
_mm_loadu_si128(reinterpret_cast<const __m128i *>(query + dim)));
__m256 data_regs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
data_regs[i] = _mm256_cvtph_ps(
_mm_loadu_si128(reinterpret_cast<const __m128i *>(ptrs[i] + dim)));
__m256 diff = _mm256_sub_ps(q, data_regs[i]);
accs[i] = _mm256_fmadd_ps(diff, diff, accs[i]);
}
dim += 8;
}
for (size_t i = 0; i < dp_batch; ++i) {
results[i] = HorizontalAdd_FP32_V256(accs[i]);
}
for (; dim < dimensionality; ++dim) {
for (size_t i = 0; i < dp_batch; ++i) {
float diff = (*(query + dim)) - (*(ptrs[i] + dim));
results[i] += diff * diff;
}
}
}
void compute_one_to_many_squared_euclidean_avx2_fp16_1(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs, size_t dim,
float *sums) {
return compute_one_to_many_squared_euclidean_avx2_fp16<ailego::Float16, 1>(
query, ptrs, prefetch_ptrs, dim, sums);
}
void compute_one_to_many_squared_euclidean_avx2_fp16_12(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs, size_t dim,
float *sums) {
return compute_one_to_many_squared_euclidean_avx2_fp16<ailego::Float16, 12>(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
} // namespace zvec::ailego::DistanceBatch

137
src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512.cc Normal file
View File

@ -0,0 +1,137 @@
// 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 <ailego/math/matrix_utility.i>
#include <ailego/utility/math_helper.h>
#include <zvec/ailego/internal/platform.h>
#include <zvec/ailego/utility/type_helper.h>
namespace zvec::ailego::DistanceBatch {
#if defined(__AVX512F__)
template <typename ValueType, size_t dp_batch>
static std::enable_if_t<std::is_same_v<ValueType, ailego::Float16>, void>
compute_one_to_many_squared_euclidean_avx512f_fp16(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, dp_batch> &prefetch_ptrs,
size_t dimensionality, float *results) {
__m512 accs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
accs[i] = _mm512_setzero_ps();
}
size_t dim = 0;
for (; dim + 32 <= dimensionality; dim += 32) {
__m512i q =
_mm512_loadu_si512(reinterpret_cast<const __m512i *>(query + dim));
__m512 q1 = _mm512_cvtph_ps(_mm512_castsi512_si256(q));
__m512 q2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(q, 1));
__m512 data_regs_1[dp_batch];
__m512 data_regs_2[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
__m512i m =
_mm512_loadu_si512(reinterpret_cast<const __m512i *>(ptrs[i] + dim));
data_regs_1[i] = _mm512_cvtph_ps(_mm512_castsi512_si256(m));
data_regs_2[i] = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(m, 1));
}
if (prefetch_ptrs[0]) {
for (size_t i = 0; i < dp_batch; ++i) {
ailego_prefetch(prefetch_ptrs[i] + dim);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
__m512 diff1 = _mm512_sub_ps(q1, data_regs_1[i]);
accs[i] = _mm512_fmadd_ps(diff1, diff1, accs[i]);
__m512 diff2 = _mm512_sub_ps(q2, data_regs_2[i]);
accs[i] = _mm512_fmadd_ps(diff2, diff2, accs[i]);
}
}
if (dim + 16 <= dimensionality) {
__m512 q = _mm512_cvtph_ps(
_mm256_loadu_si256(reinterpret_cast<const __m256i *>(query + dim)));
for (size_t i = 0; i < dp_batch; ++i) {
__m512 m = _mm512_cvtph_ps(
_mm256_loadu_si256(reinterpret_cast<const __m256i *>(ptrs[i] + dim)));
__m512 diff = _mm512_sub_ps(m, q);
accs[i] = _mm512_fmadd_ps(diff, diff, accs[i]);
}
dim += 16;
}
__m256 acc_new[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
acc_new[i] = _mm256_add_ps(
_mm512_castps512_ps256(accs[i]),
_mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(accs[i]), 1)));
}
if (dim + 8 < dimensionality) {
__m256 q = _mm256_cvtph_ps(
_mm_loadu_si128(reinterpret_cast<const __m128i *>(query + dim)));
for (size_t i = 0; i < dp_batch; ++i) {
__m256 m = _mm256_cvtph_ps(
_mm_loadu_si128(reinterpret_cast<const __m128i *>(ptrs[i] + dim)));
__m256 diff = _mm256_sub_ps(m, q);
acc_new[i] = _mm256_fmadd_ps(diff, diff, acc_new[i]);
}
dim += 8;
}
for (size_t i = 0; i < dp_batch; ++i) {
results[i] = HorizontalAdd_FP32_V256(acc_new[i]);
}
for (; dim < dimensionality; ++dim) {
for (size_t i = 0; i < dp_batch; ++i) {
float diff = (*(query + dim)) - (*(ptrs[i] + dim));
results[i] += diff * diff;
}
}
}
void compute_one_to_many_squared_euclidean_avx512f_fp16_1(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs, size_t dim,
float *sums) {
return compute_one_to_many_squared_euclidean_avx512f_fp16<ailego::Float16, 1>(
query, ptrs, prefetch_ptrs, dim, sums);
}
void compute_one_to_many_squared_euclidean_avx512f_fp16_12(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs, size_t dim,
float *sums) {
return compute_one_to_many_squared_euclidean_avx512f_fp16<ailego::Float16,
12>(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
} // namespace zvec::ailego::DistanceBatch

99
src/ailego/math_batch/euclidean_distance_batch_impl_fp16_avx512fp16.cc Normal file
View File

@ -0,0 +1,99 @@
// 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 <ailego/math/matrix_utility.i>
#include <ailego/utility/math_helper.h>
#include <zvec/ailego/internal/platform.h>
#include <zvec/ailego/utility/type_helper.h>
namespace zvec::ailego::DistanceBatch {
#if defined(__AVX512FP16__)
template <typename ValueType, size_t dp_batch>
static std::enable_if_t<std::is_same_v<ValueType, ailego::Float16>, void>
compute_one_to_many_squared_euclidean_avx512fp16_fp16(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, dp_batch> &prefetch_ptrs,
size_t dimensionality, float *results) {
__m512h accs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
accs[i] = _mm512_setzero_ph();
}
size_t dim = 0;
for (; dim + 32 <= dimensionality; dim += 32) {
__m512h q = _mm512_loadu_ph(query + dim);
__m512h data_regs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
data_regs[i] = _mm512_loadu_ph(ptrs[i] + dim);
}
if (prefetch_ptrs[0]) {
for (size_t i = 0; i < dp_batch; ++i) {
ailego_prefetch(prefetch_ptrs[i] + dim);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
__m512h diff = _mm512_sub_ph(data_regs[i], q);
accs[i] = _mm512_fmadd_ph(diff, diff, accs[i]);
}
}
if (dim < dimensionality) {
__mmask32 mask = (__mmask32)(((uint32_t)1 << (dimensionality - dim)) - 1);
for (size_t i = 0; i < dp_batch; ++i) {
__m512i zmm_undefined = _mm512_undefined_epi32();
__m512h zmm_undefined_ph = _mm512_undefined_ph();
__m512h zmm_d =
_mm512_mask_sub_ph(zmm_undefined_ph, mask,
_mm512_castsi512_ph(_mm512_mask_loadu_epi16(
zmm_undefined, mask, query + dim)),
_mm512_castsi512_ph(_mm512_mask_loadu_epi16(
zmm_undefined, mask, ptrs[i] + dim)));
accs[i] = _mm512_mask3_fmadd_ph(zmm_d, zmm_d, accs[i], mask);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
results[i] = HorizontalAdd_FP16_V512(accs[i]);
}
}
void compute_one_to_many_squared_euclidean_avx512fp16_fp16_1(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 1> &prefetch_ptrs, size_t dim,
float *sums) {
return compute_one_to_many_squared_euclidean_avx512fp16_fp16<ailego::Float16,
1>(
query, ptrs, prefetch_ptrs, dim, sums);
}
void compute_one_to_many_squared_euclidean_avx512fp16_fp16_12(
const ailego::Float16 *query, const ailego::Float16 **ptrs,
std::array<const ailego::Float16 *, 12> &prefetch_ptrs, size_t dim,
float *sums) {
return compute_one_to_many_squared_euclidean_avx512fp16_fp16<ailego::Float16,
12>(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
} // namespace zvec::ailego::DistanceBatch

104
src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx2.cc Normal file
View File

@ -0,0 +1,104 @@
// 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 <ailego/math/euclidean_distance_matrix.h>
#include <ailego/math/matrix_utility.i>
#include <ailego/utility/math_helper.h>
#include <zvec/ailego/internal/platform.h>
#include <zvec/ailego/utility/type_helper.h>
#include "distance_batch_math.h"
#define SSD_FP32_GENERAL(m, q, sum) \
{ \
float x = m - q; \
sum += (x * x); \
}
namespace zvec::ailego::DistanceBatch {
#if defined(__AVX2__)
template <typename ValueType, size_t dp_batch>
static std::enable_if_t<std::is_same_v<ValueType, float>, void>
compute_one_to_many_squared_euclidean_avx2_fp32(
const ValueType *query, const ValueType **ptrs,
std::array<const ValueType *, dp_batch> &prefetch_ptrs,
size_t dimensionality, float *results) {
__m256 accs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
accs[i] = _mm256_setzero_ps();
}
size_t dim = 0;
for (; dim + 8 <= dimensionality; dim += 8) {
__m256 q = _mm256_loadu_ps(query + dim);
__m256 data_regs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
data_regs[i] = _mm256_loadu_ps(ptrs[i] + dim);
}
if (prefetch_ptrs[0]) {
for (size_t i = 0; i < dp_batch; ++i) {
ailego_prefetch(prefetch_ptrs[i] + dim);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
__m256 diff = _mm256_sub_ps(q, data_regs[i]);
accs[i] = _mm256_fmadd_ps(diff, diff, accs[i]);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
results[i] = HorizontalAdd_FP32_V256(accs[i]);
switch (dimensionality - dim) {
case 7:
SSD_FP32_GENERAL(query[dim + 6], ptrs[i][dim + 6], results[i]);
/* FALLTHRU */
case 6:
SSD_FP32_GENERAL(query[dim + 5], ptrs[i][dim + 5], results[i]);
/* FALLTHRU */
case 5:
SSD_FP32_GENERAL(query[dim + 4], ptrs[i][dim + 4], results[i]);
/* FALLTHRU */
case 4:
SSD_FP32_GENERAL(query[dim + 3], ptrs[i][dim + 3], results[i]);
/* FALLTHRU */
case 3:
SSD_FP32_GENERAL(query[dim + 2], ptrs[i][dim + 2], results[i]);
/* FALLTHRU */
case 2:
SSD_FP32_GENERAL(query[dim + 1], ptrs[i][dim + 1], results[i]);
/* FALLTHRU */
case 1:
SSD_FP32_GENERAL(query[dim + 0], ptrs[i][dim + 0], results[i]);
}
}
}
void compute_one_to_many_squared_euclidean_avx2_fp32_1(
const float *query, const float **ptrs,
std::array<const float *, 1> &prefetch_ptrs, size_t dim, float *sums) {
return compute_one_to_many_squared_euclidean_avx2_fp32<float, 1>(
query, ptrs, prefetch_ptrs, dim, sums);
}
void compute_one_to_many_squared_euclidean_avx2_fp32_12(
const float *query, const float **ptrs,
std::array<const float *, 12> &prefetch_ptrs, size_t dim, float *sums) {
return compute_one_to_many_squared_euclidean_avx2_fp32<float, 12>(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
} // namespace zvec::ailego::DistanceBatch

95
src/ailego/math_batch/euclidean_distance_batch_impl_fp32_avx512.cc Normal file
View File

@ -0,0 +1,95 @@
// 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 <ailego/math/euclidean_distance_matrix.h>
#include <ailego/math/matrix_utility.i>
#include <ailego/utility/math_helper.h>
#include <zvec/ailego/internal/platform.h>
#include <zvec/ailego/utility/type_helper.h>
#include "distance_batch_math.h"
#define SSD_FP32_GENERAL(m, q, sum) \
{ \
float x = m - q; \
sum += (x * x); \
}
namespace zvec::ailego::DistanceBatch {
#if defined(__AVX512F__)
template <typename ValueType, size_t dp_batch>
static std::enable_if_t<std::is_same_v<ValueType, float>, void>
compute_one_to_many_squared_euclidean_avx512f_fp32(
const ValueType *query, const ValueType **ptrs,
std::array<const ValueType *, dp_batch> &prefetch_ptrs,
size_t dimensionality, float *results) {
__m512 accs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
accs[i] = _mm512_setzero_ps();
}
size_t dim = 0;
for (; dim + 16 <= dimensionality; dim += 16) {
__m512 q = _mm512_loadu_ps(query + dim);
__m512 data_regs[dp_batch];
for (size_t i = 0; i < dp_batch; ++i) {
data_regs[i] = _mm512_loadu_ps(ptrs[i] + dim);
}
if (prefetch_ptrs[0]) {
for (size_t i = 0; i < dp_batch; ++i) {
ailego_prefetch(prefetch_ptrs[i] + dim);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
__m512 diff = _mm512_sub_ps(q, data_regs[i]);
accs[i] = _mm512_fmadd_ps(diff, diff, accs[i]);
}
}
if (dim < dimensionality) {
__mmask16 mask = (__mmask16)((1 << (dimensionality - dim)) - 1);
for (size_t i = 0; i < dp_batch; ++i) {
__m512 zmm_undefined = _mm512_undefined_ps();
__m512 q = _mm512_mask_loadu_ps(zmm_undefined, mask, query + dim);
__m512 m = _mm512_mask_loadu_ps(zmm_undefined, mask, ptrs[i] + dim);
__m512 diff = _mm512_mask_sub_ps(zmm_undefined, mask, q, m);
accs[i] = _mm512_mask3_fmadd_ps(diff, diff, accs[i], mask);
}
}
for (size_t i = 0; i < dp_batch; ++i) {
results[i] = HorizontalAdd_FP32_V512(accs[i]);
}
}
void compute_one_to_many_squared_euclidean_avx512f_fp32_1(
const float *query, const float **ptrs,
std::array<const float *, 1> &prefetch_ptrs, size_t dim, float *sums) {
return compute_one_to_many_squared_euclidean_avx512f_fp32<float, 1>(
query, ptrs, prefetch_ptrs, dim, sums);
}
void compute_one_to_many_squared_euclidean_avx512f_fp32_12(
const float *query, const float **ptrs,
std::array<const float *, 12> &prefetch_ptrs, size_t dim, float *sums) {
return compute_one_to_many_squared_euclidean_avx512f_fp32<float, 12>(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
} // namespace zvec::ailego::DistanceBatch

1
src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512fp16.cc
View File

@ -168,7 +168,6 @@ void compute_one_to_many_inner_product_avx512_vnni_int8_12(
query, ptrs, prefetch_ptrs, dim, sums);
}
#endif
} // namespace zvec::ailego::DistanceBatch

48
src/core/metric/euclidean_metric.cc
View File

@ -853,37 +853,25 @@ class SquaredEuclideanMetric : public IndexMetric {
//! Retrieve distance function for query
MatrixBatchDistance batch_distance(void) const override {
switch (data_type_) {
case IndexMeta::DataType::DT_BINARY32:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::HammingDistanceMatrix, uint32_t, 1,
1>::ComputeBatch);
#if defined(AILEGO_M64)
case IndexMeta::DataType::DT_BINARY64:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::HammingDistanceMatrix, uint64_t, 1,
1>::ComputeBatch);
#endif // AILEGO_M64
case IndexMeta::DataType::DT_FP16:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix,
ailego::Float16, 1, 1>::ComputeBatch);
ailego::Float16, 12, 2>::ComputeBatch);
case IndexMeta::DataType::DT_FP32:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix, float,
1, 1>::ComputeBatch);
12, 2>::ComputeBatch);
case IndexMeta::DataType::DT_INT8:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix, int8_t,
1, 1>::ComputeBatch);
12, 2>::ComputeBatch);
case IndexMeta::DataType::DT_INT4:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::SquaredEuclideanDistanceMatrix,
uint8_t, 1, 1>::ComputeBatch);
uint8_t, 12, 2>::ComputeBatch);
default:
return nullptr;
@ -1009,6 +997,34 @@ class EuclideanMetric : public IndexMetric {
}
}
//! Retrieve distance function for query
MatrixBatchDistance batch_distance(void) const override {
switch (data_type_) {
case IndexMeta::DataType::DT_FP16:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::EuclideanDistanceMatrix,
ailego::Float16, 12, 2>::ComputeBatch);
case IndexMeta::DataType::DT_FP32:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::EuclideanDistanceMatrix, float, 12,
2>::ComputeBatch);
case IndexMeta::DataType::DT_INT8:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::EuclideanDistanceMatrix, int8_t, 12,
2>::ComputeBatch);
case IndexMeta::DataType::DT_INT4:
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::EuclideanDistanceMatrix, uint8_t, 12,
2>::ComputeBatch);
default:
return nullptr;
}
}
//! Retrieve params of Metric
const ailego::Params &params(void) const override {
return params_;

16
src/core/metric/hamming_metric.cc
View File

@ -189,22 +189,6 @@ class HammingMetric : public IndexMetric {
return nullptr;
}
MatrixBatchDistance batch_distance(void) const override {
#if defined(AILEGO_M64)
if (feature_type_ == IndexMeta::DataType::DT_BINARY64) {
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::HammingDistanceMatrix, uint64_t, 1,
1>::ComputeBatch);
}
#endif
if (feature_type_ == IndexMeta::DataType::DT_BINARY32) {
return reinterpret_cast<IndexMetric::MatrixBatchDistanceHandle>(
ailego::BaseDistance<ailego::HammingDistanceMatrix, uint32_t, 1,
1>::ComputeBatch);
}
return nullptr;
}
//! Retrieve distance function for index features
MatrixDistance distance_matrix(size_t m, size_t n) const override {
#if defined(AILEGO_M64)

5
tests/core/algorithm/hnsw/hnsw_streamer_test.cc
View File

@ -3468,6 +3468,7 @@ TEST_F(HnswStreamerTest, TestGroupInBruteforceSearch) {
}
}
#if 0
TEST_F(HnswStreamerTest, TestBinaryConverter) {
uint32_t dimension = 2560;
@ -3552,6 +3553,7 @@ TEST_F(HnswStreamerTest, TestBinaryConverter) {
ASSERT_NEAR(0, results[0].score(), epison);
}
}
#endif
TEST_F(HnswStreamerTest, TestAddAndSearchWithID) {
IndexStreamer::Pointer streamer =
@ -3666,6 +3668,7 @@ TEST_F(HnswStreamerTest, TestAddAndSearchWithID) {
// EXPECT_GT(cost, 2.0f);
}
#if 0
TEST_F(HnswStreamerTest, TestBasicRefiner) {
uint32_t dimension = 1120;
@ -3788,6 +3791,8 @@ TEST_F(HnswStreamerTest, TestBasicRefiner) {
}
}
#endif
} // namespace core
} // namespace zvec