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

fix/fix mips euclidean (#226) 

* fix: fix mips euclidean
This commit is contained in:
rayx 2026-03-16 13:41:58 +08:00 committed by GitHub
parent 65b45470d6
commit de1e7050d7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
20 changed files with 575 additions and 207 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/ailego/math/inner_product_matrix_fp32_avx.cc
View File

@ -88,6 +88,10 @@ float InnerProductAVX(const float *lhs, const float *rhs, size_t size) {
return result;
}
float MinusInnerProductAVX(const float *lhs, const float *rhs, size_t size) {
return -1 * InnerProductAVX(lhs, rhs, size);
}
#endif // __AVX__
} // namespace ailego

4
src/ailego/math/inner_product_matrix_fp32_avx512.cc
View File

@ -69,6 +69,10 @@ float InnerProductAVX512(const float *lhs, const float *rhs, size_t size) {
return HorizontalAdd_FP32_V512(zmm_sum_0);
}
float MinusInnerProductAVX512(const float *lhs, const float *rhs, size_t size) {
return -1 * InnerProductAVX512(lhs, rhs, size);
}
#endif
} // namespace ailego

9
src/ailego/math/inner_product_matrix_fp32_dispatch.cc
View File

@ -25,6 +25,7 @@ float MinusInnerProductNEON(const float *lhs, const float *rhs, size_t size);
#if defined(__AVX512F__)
float InnerProductAVX512(const float *lhs, const float *rhs, size_t size);
float MinusInnerProductAVX512(const float *lhs, const float *rhs, size_t size);
#endif
#if defined(__AVX__)
@ -70,12 +71,12 @@ void MinusInnerProductMatrix<float, 1, 1>::Compute(const ValueType *m,
const ValueType *q,
size_t dim, float *out) {
#if defined(__ARM_NEON)
*out = -InnerProductNEON(m, q, dim);
*out = MinusInnerProductNEON(m, q, dim);
#else
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
if (dim > 15) {
*out = -InnerProductAVX512(m, q, dim);
*out = MinusInnerProductAVX512(m, q, dim);
return;
}
}
@ -83,12 +84,12 @@ void MinusInnerProductMatrix<float, 1, 1>::Compute(const ValueType *m,
#if defined(__AVX__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {
if (dim > 7) {
*out = -InnerProductAVX(m, q, dim);
*out = MinusInnerProductAVX(m, q, dim);
return;
}
}
#endif // __AVX__
*out = -InnerProductSSE(m, q, dim);
*out = MinusInnerProductSSE(m, q, dim);
#endif // __ARM_NEON
}

5
src/ailego/math/inner_product_matrix_fp32_neon.cc
View File

@ -51,6 +51,11 @@ float InnerProductNEON(const float *lhs, const float *rhs, size_t size) {
}
return result;
}
float MinusInnerProductNEON(const float *lhs, const float *rhs, size_t size) {
return -1 * InnerProductNEON(lhs, rhs, size);
}
#endif // __ARM_NEON
} // namespace ailego

5
src/ailego/math/inner_product_matrix_fp32_sse.cc
View File

@ -74,6 +74,11 @@ float InnerProductSSE(const float *lhs, const float *rhs, size_t size) {
return result;
}
float MinusInnerProductSSE(const float *lhs, const float *rhs, size_t size) {
return -1 * InnerProductSSE(lhs, rhs, size);
}
#endif // __SSE__
// #if 1

37
src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc
View File

@ -110,7 +110,42 @@ float InnerProductAndSquaredNormAVX(const Float16 *lhs, const Float16 *rhs,
*sqr = norm2;
return result;
}
float MipsEucldeanDistanceSphericalInjectionAVX(const Float16 *lhs,
const Float16 *rhs, size_t size,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX(lhs, rhs, size, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX(const Float16 *lhs,
const Float16 *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX(lhs, rhs, size, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __AVX__ && __F16C__
} // namespace ailego
} // namespace zvec
} // namespace zvec

37
src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc
View File

@ -128,7 +128,42 @@ float InnerProductAndSquaredNormAVX512(const Float16 *lhs, const Float16 *rhs,
*sqr = norm2;
return result;
}
float MipsEucldeanDistanceSphericalInjectionAVX512(const Float16 *lhs,
const Float16 *rhs,
size_t size, float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX512(lhs, rhs, size, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX512(const Float16 *lhs,
const Float16 *rhs,
size_t size,
size_t m, float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX512(lhs, rhs, size, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __AVX512F__
} // namespace ailego
} // namespace zvec
} // namespace zvec

76
src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc
View File

@ -19,18 +19,33 @@ namespace zvec {
namespace ailego {
#if defined(__ARM_NEON)
float InnerProductAndSquaredNormNEON(const Float16 *lhs, const Float16 *rhs,
size_t size, float *sql, float *sqr);
float MipsEucldeanDistanceRepeatedQuadraticInjectionNEON(const Float16 *lhs,
const Float16 *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionNEON(const Float16 *lhs,
const Float16 *rhs,
size_t size, float e2);
#endif
#if defined(__AVX512F__)
float InnerProductAndSquaredNormAVX512(const Float16 *lhs, const Float16 *rhs,
size_t size, float *sql, float *sqr);
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX512(const Float16 *lhs,
const Float16 *rhs,
size_t size,
size_t m, float e2);
float MipsEucldeanDistanceSphericalInjectionAVX512(const Float16 *lhs,
const Float16 *rhs,
size_t size, float e2);
#endif
#if defined(__AVX__)
float InnerProductAndSquaredNormAVX(const Float16 *lhs, const Float16 *rhs,
size_t size, float *sql, float *sqr);
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX(const Float16 *lhs,
const Float16 *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionAVX(const Float16 *lhs,
const Float16 *rhs, size_t size,
float e2);
#endif
#if (defined(__F16C__) && defined(__AVX__)) || \
@ -38,59 +53,38 @@ float InnerProductAndSquaredNormAVX(const Float16 *lhs, const Float16 *rhs,
//! Compute the distance between matrix and query by SphericalInjection
void MipsSquaredEuclideanDistanceMatrix<Float16, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__ARM_NEON)
sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2);
*out = MipsEucldeanDistanceSphericalInjectionNEON(p, q, dim, e2);
#else
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2);
} else
#endif //__AVX512F__
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {
sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2);
}
*out = MipsEucldeanDistanceSphericalInjectionAVX512(p, q, dim, e2);
return;
}
#endif
*out = MipsEucldeanDistanceSphericalInjectionAVX(p, q, dim, e2);
#endif //__ARM_NEON
*out = ComputeSphericalInjection(sum, u2, v2, e2);
}
//! Compute the distance between matrix and query by RepeatedQuadraticInjection
void MipsSquaredEuclideanDistanceMatrix<Float16, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,
float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__ARM_NEON)
sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2);
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionNEON(p, q, dim, m, e2);
#else
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2);
} else
#endif //__AVX512F__
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {
sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2);
}
#endif //__ARM_NEON
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
*out =
MipsEucldeanDistanceRepeatedQuadraticInjectionAVX512(p, q, dim, m, e2);
return;
}
*out = sum;
#endif
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionAVX(p, q, dim, m, e2);
#endif //__ARM_NEON
}
#endif // (__F16C__ && __AVX__) || (__ARM_NEON && __aarch64__)
} // namespace ailego
} // namespace zvec
} // namespace zvec

37
src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc
View File

@ -119,8 +119,43 @@ float InnerProductAndSquaredNormNEON(const Float16 *lhs, const Float16 *rhs,
*sqr = norm2;
return result;
}
#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
float MipsEucldeanDistanceSphericalInjectionNEON(const Float16 *lhs,
const Float16 *rhs,
size_t size, float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormNEON(lhs, rhs, size, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionNEON(const Float16 *lhs,
const Float16 *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormNEON(lhs, rhs, size, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __ARM_NEON && __aarch64__
} // namespace ailego
} // namespace zvec
} // namespace zvec

49
src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc
View File

@ -19,6 +19,11 @@
namespace zvec {
namespace ailego {
#if defined(__SSE__)
float InnerProductAndSquaredNormSSE(const float *lhs, const float *rhs,
size_t size, float *sql, float *sqr);
#endif
#if defined(__AVX__)
//! Compute the Inner Product between p and q, and each Squared L2-Norm value
float InnerProductAndSquaredNormAVX(const float *lhs, const float *rhs,
@ -108,7 +113,49 @@ float InnerProductAndSquaredNormAVX(const float *lhs, const float *rhs,
*sqr = norm2;
return result;
}
float MipsEucldeanDistanceSphericalInjectionAVX(const float *lhs,
const float *rhs, size_t size,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
if (size > 7) {
sum = InnerProductAndSquaredNormAVX(lhs, rhs, size, &u2, &v2);
} else {
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
}
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX(const float *lhs,
const float *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
if (size > 7) {
sum = InnerProductAndSquaredNormAVX(lhs, rhs, size, &u2, &v2);
} else {
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
}
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __AVX__
} // namespace ailego
} // namespace zvec
} // namespace zvec

58
src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc
View File

@ -19,6 +19,16 @@
namespace zvec {
namespace ailego {
#if defined(__SSE__)
float InnerProductAndSquaredNormSSE(const float *lhs, const float *rhs,
size_t size, float *sql, float *sqr);
#endif
#if defined(__AVX__)
float InnerProductAndSquaredNormAVX(const float *lhs, const float *rhs,
size_t size, float *sql, float *sqr);
#endif
#if defined(__AVX512F__)
//! Compute the Inner Product between p and q, and each Squared L2-Norm value
float InnerProductAndSquaredNormAVX512(const float *lhs, const float *rhs,
@ -94,7 +104,53 @@ float InnerProductAndSquaredNormAVX512(const float *lhs, const float *rhs,
*sqr = HorizontalAdd_FP32_V512(zmm_sum_norm2);
return HorizontalAdd_FP32_V512(zmm_sum_0);
}
float MipsEucldeanDistanceSphericalInjectionAVX512(const float *lhs,
const float *rhs,
size_t size, float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
if (size > 15) {
sum = InnerProductAndSquaredNormAVX512(lhs, rhs, size, &u2, &v2);
} else if (size > 7) {
sum = InnerProductAndSquaredNormAVX(lhs, rhs, size, &u2, &v2);
} else {
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
}
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX512(const float *lhs,
const float *rhs,
size_t size,
size_t m, float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
if (size > 15) {
sum = InnerProductAndSquaredNormAVX512(lhs, rhs, size, &u2, &v2);
} else if (size > 7) {
sum = InnerProductAndSquaredNormAVX(lhs, rhs, size, &u2, &v2);
} else {
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
}
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __AVX512F__
} // namespace ailego
} // namespace zvec
} // namespace zvec

123
src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc
View File

@ -24,18 +24,33 @@ float InnerProductAndSquaredNormNEON(const float *lhs, const float *rhs,
#endif
#if defined(__AVX512F__)
float InnerProductAndSquaredNormAVX512(const float *lhs, const float *rhs,
size_t size, float *sql, float *sqr);
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX512(const float *lhs,
const float *rhs,
size_t size,
size_t m, float e2);
float MipsEucldeanDistanceSphericalInjectionAVX512(const float *lhs,
const float *rhs,
size_t size, float e2);
#endif
#if defined(__AVX__)
float InnerProductAndSquaredNormAVX(const float *lhs, const float *rhs,
size_t size, float *sql, float *sqr);
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX(const float *lhs,
const float *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionAVX(const float *lhs,
const float *rhs, size_t size,
float e2);
#endif
#if defined(__SSE__)
float InnerProductAndSquaredNormSSE(const float *lhs, const float *rhs,
size_t size, float *sql, float *sqr);
float MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(const float *lhs,
const float *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionSSE(const float *lhs,
const float *rhs, size_t size,
float e2);
#endif
#if defined(__SSE4_1__)
@ -58,58 +73,39 @@ float MipsInnerProductSparseInSegment(uint32_t m_sparse_count,
//! Compute the distance between matrix and query by SphericalInjection
void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && dim > 15) {
sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2);
} else
#endif // __AVX512F__
#if defined(__AVX__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX && dim > 7) {
sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2);
} else
#endif // __AVX__
{
sum = InnerProductAndSquaredNormSSE(p, q, dim, &u2, &v2);
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
*out = MipsEucldeanDistanceSphericalInjectionAVX512(p, q, dim, e2);
return;
}
*out = ComputeSphericalInjection(sum, u2, v2, e2);
#endif //__AVX512F__
#if defined(__AVX__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {
*out = MipsEucldeanDistanceSphericalInjectionAVX(p, q, dim, e2);
return;
}
#endif // __AVX__
*out = MipsEucldeanDistanceSphericalInjectionSSE(p, q, dim, e2);
}
//! Compute the distance between matrix and query by RepeatedQuadraticInjection
void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,
float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__AVX512F__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && dim > 15) {
sum = InnerProductAndSquaredNormAVX512(p, q, dim, &u2, &v2);
} else
#endif // __AVX512F__
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
*out =
MipsEucldeanDistanceRepeatedQuadraticInjectionAVX512(p, q, dim, m, e2);
return;
}
#endif //__AVX512F__
#if defined(__AVX__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX && dim > 7) {
sum = InnerProductAndSquaredNormAVX(p, q, dim, &u2, &v2);
} else
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionAVX(p, q, dim, m, e2);
return;
}
#endif // __AVX__
{
sum = InnerProductAndSquaredNormSSE(p, q, dim, &u2, &v2);
}
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
*out = sum;
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(p, q, dim, m, e2);
}
#endif // __SSE__
@ -132,5 +128,36 @@ float MipsSquaredEuclideanSparseDistanceMatrix<float>::
#endif
}
#if defined(__ARM_NEON)
//! Compute the distance between matrix and query by SphericalInjection
void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {
float u2{0.0f};
float v2{0.0f};
float sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2);
*out = ComputeSphericalInjection(sum, u2, v2, e2);
}
//! Compute the distance between matrix and query by RepeatedQuadraticInjection
void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,
float *out) {
float u2{0.0f};
float v2{0.0f};
float sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
*out = sum;
}
#endif //__ARM_NEON
} // namespace ailego
} // namespace zvec
} // namespace zvec

30
src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc
View File

@ -71,35 +71,7 @@ float InnerProductAndSquaredNormNEON(const float *lhs, const float *rhs,
return result;
}
//! Compute the distance between matrix and query by SphericalInjection
void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {
float u2;
float v2;
float sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2);
*out = ComputeSphericalInjection(sum, u2, v2, e2);
}
//! Compute the distance between matrix and query by RepeatedQuadraticInjection
void MipsSquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,
float *out) {
float u2;
float v2;
float sum = InnerProductAndSquaredNormNEON(p, q, dim, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
*out = sum;
}
#endif //__ARM_NEON
} // namespace ailego
} // namespace zvec
} // namespace zvec

36
src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc
View File

@ -96,6 +96,40 @@ float InnerProductAndSquaredNormSSE(const float *lhs, const float *rhs,
return result;
}
float MipsEucldeanDistanceSphericalInjectionSSE(const float *lhs,
const float *rhs, size_t size,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(const float *lhs,
const float *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __SSE__
// #if 1
@ -333,4 +367,4 @@ float MipsInnerProductSparseInSegment(uint32_t m_sparse_count,
#endif // __SSE4_1__
} // namespace ailego
} // namespace zvec
} // namespace zvec

40
src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc
View File

@ -23,8 +23,8 @@ namespace ailego {
#if defined(__AVX2__)
//! Compute the Inner Product between p and q, and each Squared L2-Norm value
float InnerProductAndSquaredNormAVX(const uint8_t *lhs, const uint8_t *rhs,
size_t size, float *sql, float *sqr) {
float InnerProductAndSquaredNormAVX2(const uint8_t *lhs, const uint8_t *rhs,
size_t size, float *sql, float *sqr) {
const uint8_t *last = lhs + size;
const uint8_t *last_aligned = lhs + ((size >> 5) << 5);
__m256i ymm_sum_0 = _mm256_setzero_si256();
@ -134,7 +134,41 @@ float InnerProductAndSquaredNormAVX(const uint8_t *lhs, const uint8_t *rhs,
*sqr = norm2;
return result;
}
float MipsEucldeanDistanceSphericalInjectionAVX2(const uint8_t *lhs,
const uint8_t *rhs,
size_t size, float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX2(lhs, rhs, size >> 1, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX2(const uint8_t *lhs,
const uint8_t *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX2(lhs, rhs, size >> 1, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __AVX2__
} // namespace ailego
} // namespace zvec
} // namespace zvec

62
src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc
View File

@ -20,64 +20,52 @@
namespace zvec {
namespace ailego {
#if defined(__AVX__)
float InnerProductAndSquaredNormAVX(const uint8_t *lhs, const uint8_t *rhs,
size_t size, float *sql, float *sqr);
#if defined(__AVX2__)
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX2(const uint8_t *lhs,
const uint8_t *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionAVX2(const uint8_t *lhs,
const uint8_t *rhs,
size_t size, float e2);
#endif
#if defined(__SSE__)
float InnerProductAndSquaredNormSSE(const uint8_t *lhs, const uint8_t *rhs,
size_t size, float *sql, float *sqr);
#if defined(__SSE4_1__)
float MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(const uint8_t *lhs,
const uint8_t *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionSSE(const uint8_t *lhs,
const uint8_t *rhs, size_t size,
float e2);
#endif
#if defined(__SSE4_1__)
//! Compute the distance between matrix and query by SphericalInjection
void MipsSquaredEuclideanDistanceMatrix<uint8_t, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
sum = InnerProductAndSquaredNormAVX(p, q, dim >> 1, &u2, &v2);
} else
#endif
{
sum = InnerProductAndSquaredNormSSE(p, q, dim >> 1, &u2, &v2);
*out = MipsEucldeanDistanceSphericalInjectionAVX2(p, q, dim, e2);
return;
}
*out = ComputeSphericalInjection(sum, u2, v2, e2);
#endif
*out = MipsEucldeanDistanceSphericalInjectionSSE(p, q, dim, e2);
}
//! Compute the distance between matrix and query by RepeatedQuadraticInjection
void MipsSquaredEuclideanDistanceMatrix<uint8_t, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,
float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
sum = InnerProductAndSquaredNormAVX(p, q, dim >> 1, &u2, &v2);
} else
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionAVX2(p, q, dim, m, e2);
return;
}
#endif
{
sum = InnerProductAndSquaredNormSSE(p, q, dim >> 1, &u2, &v2);
}
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
*out = sum;
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(p, q, dim, m, e2);
}
#endif
} // namespace ailego
} // namespace zvec
} // namespace zvec

37
src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc
View File

@ -98,7 +98,42 @@ float InnerProductAndSquaredNormSSE(const uint8_t *lhs, const uint8_t *rhs,
*sqr = norm2;
return result;
}
float MipsEucldeanDistanceSphericalInjectionSSE(const uint8_t *lhs,
const uint8_t *rhs, size_t size,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size >> 1, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(const uint8_t *lhs,
const uint8_t *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size >> 1, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __SSE4_1__
} // namespace ailego
} // namespace zvec
} // namespace zvec

36
src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc
View File

@ -153,7 +153,41 @@ float InnerProductAndSquaredNormAVX2(const int8_t *lhs, const int8_t *rhs,
*sqr = norm2;
return result;
}
float MipsEucldeanDistanceSphericalInjectionAVX2(const int8_t *lhs,
const int8_t *rhs, size_t size,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX2(lhs, rhs, size, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX2(const int8_t *lhs,
const int8_t *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormAVX2(lhs, rhs, size, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __AVX2__
} // namespace ailego
} // namespace zvec
} // namespace zvec

60
src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc
View File

@ -19,63 +19,51 @@ namespace zvec {
namespace ailego {
#if defined(__AVX2__)
float InnerProductAndSquaredNormAVX2(const int8_t *lhs, const int8_t *rhs,
size_t size, float *sql, float *sqr);
float MipsEucldeanDistanceRepeatedQuadraticInjectionAVX2(const int8_t *lhs,
const int8_t *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionAVX2(const int8_t *lhs,
const int8_t *rhs, size_t size,
float e2);
#endif
#if defined(__SSE__)
float InnerProductAndSquaredNormSSE(const int8_t *lhs, const int8_t *rhs,
size_t size, float *sql, float *sqr);
#if defined(__SSE4_1__)
float MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(const int8_t *lhs,
const int8_t *rhs,
size_t size, size_t m,
float e2);
float MipsEucldeanDistanceSphericalInjectionSSE(const int8_t *lhs,
const int8_t *rhs, size_t size,
float e2);
#endif
#if defined(__SSE4_1__)
//! Compute the distance between matrix and query by SphericalInjection
void MipsSquaredEuclideanDistanceMatrix<int8_t, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, float e2, float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
sum = InnerProductAndSquaredNormAVX2(p, q, dim, &u2, &v2);
} else
#endif
{
sum = InnerProductAndSquaredNormSSE(p, q, dim, &u2, &v2);
*out = MipsEucldeanDistanceSphericalInjectionAVX2(p, q, dim, e2);
return;
}
*out = ComputeSphericalInjection(sum, u2, v2, e2);
#endif
*out = MipsEucldeanDistanceSphericalInjectionSSE(p, q, dim, e2);
}
//! Compute the distance between matrix and query by RepeatedQuadraticInjection
void MipsSquaredEuclideanDistanceMatrix<int8_t, 1, 1>::Compute(
const ValueType *p, const ValueType *q, size_t dim, size_t m, float e2,
float *out) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
#if defined(__AVX2__)
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2) {
sum = InnerProductAndSquaredNormAVX2(p, q, dim, &u2, &v2);
} else
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionAVX2(p, q, dim, m, e2);
return;
}
#endif
{
sum = InnerProductAndSquaredNormSSE(p, q, dim, &u2, &v2);
}
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
*out = sum;
*out = MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(p, q, dim, m, e2);
}
#endif // __SSE4_1__
} // namespace ailego
} // namespace zvec
} // namespace zvec

37
src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc
View File

@ -131,7 +131,42 @@ float InnerProductAndSquaredNormSSE(const int8_t *lhs, const int8_t *rhs,
*sqr = norm2;
return result;
}
float MipsEucldeanDistanceSphericalInjectionSSE(const int8_t *lhs,
const int8_t *rhs, size_t size,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
return ComputeSphericalInjection(sum, u2, v2, e2);
}
float MipsEucldeanDistanceRepeatedQuadraticInjectionSSE(const int8_t *lhs,
const int8_t *rhs,
size_t size, size_t m,
float e2) {
float u2{0.0f};
float v2{0.0f};
float sum{0.0f};
sum = InnerProductAndSquaredNormSSE(lhs, rhs, size, &u2, &v2);
sum = e2 * (u2 + v2 - 2 * sum);
u2 *= e2;
v2 *= e2;
for (size_t i = 0; i < m; ++i) {
sum += (u2 - v2) * (u2 - v2);
u2 = u2 * u2;
v2 = v2 * v2;
}
return sum;
}
#endif // __SSE4_1__
} // namespace ailego
} // namespace zvec
} // namespace zvec