diff --git a/src/ailego/math/inner_product_matrix_fp32_avx.cc b/src/ailego/math/inner_product_matrix_fp32_avx.cc index 128adfd..23c1f13 100644 --- a/src/ailego/math/inner_product_matrix_fp32_avx.cc +++ b/src/ailego/math/inner_product_matrix_fp32_avx.cc @@ -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 diff --git a/src/ailego/math/inner_product_matrix_fp32_avx512.cc b/src/ailego/math/inner_product_matrix_fp32_avx512.cc index af3bf74..c888115 100644 --- a/src/ailego/math/inner_product_matrix_fp32_avx512.cc +++ b/src/ailego/math/inner_product_matrix_fp32_avx512.cc @@ -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 diff --git a/src/ailego/math/inner_product_matrix_fp32_dispatch.cc b/src/ailego/math/inner_product_matrix_fp32_dispatch.cc index 57acef2..175dbf9 100644 --- a/src/ailego/math/inner_product_matrix_fp32_dispatch.cc +++ b/src/ailego/math/inner_product_matrix_fp32_dispatch.cc @@ -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::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::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 } diff --git a/src/ailego/math/inner_product_matrix_fp32_neon.cc b/src/ailego/math/inner_product_matrix_fp32_neon.cc index e8626a3..011f908 100644 --- a/src/ailego/math/inner_product_matrix_fp32_neon.cc +++ b/src/ailego/math/inner_product_matrix_fp32_neon.cc @@ -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 diff --git a/src/ailego/math/inner_product_matrix_fp32_sse.cc b/src/ailego/math/inner_product_matrix_fp32_sse.cc index 8a302bf..f90801e 100644 --- a/src/ailego/math/inner_product_matrix_fp32_sse.cc +++ b/src/ailego/math/inner_product_matrix_fp32_sse.cc @@ -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 diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc index c93edc1..bc066ef 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc index 51ce4fc..fb87aa6 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_avx512.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc index b99ab45..be997fb 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_dispatch.cc @@ -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::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::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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc index 22493b4..8a1dd0e 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp16_neon.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc index cff60e8..ac958e8 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc index 1ac56a2..d48080e 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_avx512.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc index 992da0d..10cfec9 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_dispatch.cc @@ -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::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::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:: #endif } +#if defined(__ARM_NEON) +//! Compute the distance between matrix and query by SphericalInjection +void MipsSquaredEuclideanDistanceMatrix::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::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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc index 8e98922..ca536c3 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_neon.cc @@ -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::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::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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc b/src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc index 43d8f9b..357703d 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_fp32_sse.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc b/src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc index 95a3f00..378fd75 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_int4_avx2.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc index c967f83..238eb46 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_int4_dispatch.cc @@ -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::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::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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc b/src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc index 139b14c..0537d34 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_int4_sse.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc b/src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc index 0b96953..65a7cc8 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_int8_avx2.cc @@ -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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc b/src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc index 35b4a8c..5512c6c 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_int8_dispatch.cc @@ -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::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::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 \ No newline at end of file +} // namespace zvec diff --git a/src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc b/src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc index a0d6192..8a92f52 100644 --- a/src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc +++ b/src/ailego/math/mips_euclidean_distance_matrix_int8_sse.cc @@ -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 \ No newline at end of file +} // namespace zvec