Skip to content

Commit 2f0b573

Browse files
ihb2032ihb2032
committed
feat: add RVV non-batch distance operators
Signed-off-by: ihb2032 <hebome@foxmail.com>
1 parent 3d6906c commit 2f0b573

33 files changed

Lines changed: 1751 additions & 34 deletions

‎src/ailego/math/euclidean_distance_matrix.h‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -480,8 +480,6 @@ struct EuclideanDistanceMatrix<float, 1, 1> {
480480
static void Compute(const ValueType *m, const ValueType *q, size_t dim,
481481
float *out);
482482
};
483-
484-
485483
//--------------------------------------------------
486484
// Sparse
487485
//--------------------------------------------------

‎src/ailego/math/euclidean_distance_matrix_fp16_dispatch.cc‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,12 @@ float SquaredEuclideanDistanceFp16NEON(const Float16 *lhs, const Float16 *rhs,
2323
size_t size);
2424
#endif
2525

26+
#if defined(__riscv_zvfh)
27+
float SquaredEuclideanDistanceRVV(const Float16 *lhs, const Float16 *rhs,
28+
size_t size);
29+
float EuclideanDistanceRVV(const Float16 *lhs, const Float16 *rhs, size_t size);
30+
#endif
31+
2632
#if defined(__AVX512FP16__)
2733
float SquaredEuclideanDistanceFp16AVX512FP16(const Float16 *lhs,
2834
const Float16 *rhs, size_t size);
@@ -38,14 +44,25 @@ float SquaredEuclideanDistanceFp16AVX(const Float16 *lhs, const Float16 *rhs,
3844
size_t size);
3945
#endif
4046

47+
<<<<<<< HEAD
4148
float SquaredEuclideanDistanceFp16Scalar(const Float16 *lhs, const Float16 *rhs,
4249
size_t size);
4350

51+
=======
52+
#if (defined(__F16C__) && defined(__AVX__)) || \
53+
(defined(__ARM_NEON) && defined(__aarch64__)) || defined(__riscv_zvfh)
54+
>>>>>>> 05b06b8 (feat: add RVV non-batch distance operators)
4455
//! Compute the distance between matrix and query (FP16, M=1, N=1)
4556
void SquaredEuclideanDistanceMatrix<Float16, 1, 1>::Compute(const ValueType *m,
4657
const ValueType *q,
4758
size_t dim,
4859
float *out) {
60+
#if defined(__riscv_zvfh)
61+
if (zvec::ailego::internal::CpuFeatures::static_flags_.RISCV_ZVFH) {
62+
*out = SquaredEuclideanDistanceRVV(m, q, dim);
63+
return;
64+
}
65+
#endif
4966
#if defined(__ARM_NEON)
5067
*out = SquaredEuclideanDistanceFp16NEON(m, q, dim);
5168
#else
@@ -61,6 +78,7 @@ void SquaredEuclideanDistanceMatrix<Float16, 1, 1>::Compute(const ValueType *m,
6178
return;
6279
}
6380
#endif
81+
<<<<<<< HEAD
6482

6583
#if defined(__AVX__)
6684
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX) {
@@ -70,16 +88,42 @@ void SquaredEuclideanDistanceMatrix<Float16, 1, 1>::Compute(const ValueType *m,
7088
#endif
7189
*out = SquaredEuclideanDistanceFp16Scalar(m, q, dim);
7290

91+
=======
92+
#if defined(__AVX__)
93+
SquaredEuclideanDistanceAVX(m, q, dim, out);
94+
// ACCUM_FP16_1X1_AVX(m, q, dim, out, 0ull, )
95+
#else
96+
float sum = 0.0f;
97+
for (size_t i = 0; i < dim; ++i) {
98+
sum += MathHelper::SquaredDifference(m[i], q[i]);
99+
}
100+
*out = sum;
101+
#endif
102+
>>>>>>> 05b06b8 (feat: add RVV non-batch distance operators)
73103
#endif //__ARM_NEON
74104
}
75105

76106
//! Compute the distance between matrix and query (FP16, M=1, N=1)
77107
void EuclideanDistanceMatrix<Float16, 1, 1>::Compute(const ValueType *m,
78108
const ValueType *q,
79109
size_t dim, float *out) {
110+
#if defined(__riscv_zvfh)
111+
if (zvec::ailego::internal::CpuFeatures::static_flags_.RISCV_ZVFH) {
112+
*out = EuclideanDistanceRVV(m, q, dim);
113+
return;
114+
}
115+
#if !defined(__ARM_NEON) && !defined(__AVX__)
116+
float sum = 0.0f;
117+
for (size_t i = 0; i < dim; ++i) {
118+
sum += MathHelper::SquaredDifference(m[i], q[i]);
119+
}
120+
*out = std::sqrt(sum);
121+
return;
122+
#endif
123+
#endif
80124
SquaredEuclideanDistanceMatrix<Float16, 1, 1>::Compute(m, q, dim, out);
81125
*out = std::sqrt(*out);
82126
}
83127

84128
} // namespace ailego
85-
} // namespace zvec
129+
} // namespace zvec

‎src/ailego/math/euclidean_distance_matrix_fp16_rvv.cc‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Copyright 2025-present the zvec project
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include <cmath>
16+
#include <zvec/ailego/internal/platform.h>
17+
#include "euclidean_distance_matrix.h"
18+
19+
namespace zvec {
20+
namespace ailego {
21+
22+
#if defined(__riscv_zvfh)
23+
namespace {
24+
25+
static inline float SquaredEuclideanDistanceRVVImpl(const Float16 *lhs,
26+
const Float16 *rhs,
27+
size_t size) {
28+
const _Float16 *lhs_fp16 = reinterpret_cast<const _Float16 *>(lhs);
29+
const _Float16 *rhs_fp16 = reinterpret_cast<const _Float16 *>(rhs);
30+
const size_t vlmax = __riscv_vsetvlmax_e16m4();
31+
vfloat32m8_t v_sum = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
32+
33+
while (size != 0) {
34+
const size_t vl = __riscv_vsetvl_e16m4(size);
35+
vfloat16m4_t v_lhs = __riscv_vle16_v_f16m4(lhs_fp16, vl);
36+
vfloat16m4_t v_rhs = __riscv_vle16_v_f16m4(rhs_fp16, vl);
37+
vfloat32m8_t v_diff = __riscv_vfwsub_vv_f32m8(v_lhs, v_rhs, vl);
38+
v_sum = __riscv_vfmacc_vv_f32m8_tu(v_sum, v_diff, v_diff, vl);
39+
lhs_fp16 += vl;
40+
rhs_fp16 += vl;
41+
size -= vl;
42+
}
43+
44+
vfloat32m1_t v_zero = __riscv_vfmv_v_f_f32m1(0.0f, 1);
45+
vfloat32m1_t v_red = __riscv_vfredusum_vs_f32m8_f32m1(v_sum, v_zero, vlmax);
46+
return __riscv_vfmv_f_s_f32m1_f32(v_red);
47+
}
48+
49+
} // namespace
50+
51+
float SquaredEuclideanDistanceRVV(const Float16 *lhs, const Float16 *rhs,
52+
size_t size) {
53+
return SquaredEuclideanDistanceRVVImpl(lhs, rhs, size);
54+
}
55+
56+
float EuclideanDistanceRVV(const Float16 *lhs, const Float16 *rhs,
57+
size_t size) {
58+
return std::sqrt(SquaredEuclideanDistanceRVVImpl(lhs, rhs, size));
59+
}
60+
61+
#endif // __riscv_zvfh
62+
63+
} // namespace ailego
64+
} // namespace zvec

‎src/ailego/math/euclidean_distance_matrix_fp32_dispatch.cc‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,12 @@
1818
namespace zvec {
1919
namespace ailego {
2020

21+
#if defined(__riscv_vector)
22+
float SquaredEuclideanDistanceRVV(const float *lhs, const float *rhs,
23+
size_t size);
24+
float EuclideanDistanceRVV(const float *lhs, const float *rhs, size_t size);
25+
#endif
26+
2127
#if defined(__ARM_NEON)
2228
void SquaredEuclideanDistanceFp32NEON(const float *lhs, const float *rhs,
2329
size_t size, float *out);
@@ -44,14 +50,29 @@ float SquaredEuclideanDistanceFp32Scalar(const float *lhs, const float *rhs,
4450
//-----------------------------------------------------------
4551
// SquaredEuclideanDistance
4652
//-----------------------------------------------------------
53+
<<<<<<< HEAD
54+
=======
55+
#if defined(__SSE__) || defined(__ARM_NEON) || defined(__riscv_vector)
56+
>>>>>>> 05b06b8 (feat: add RVV non-batch distance operators)
4757
//! Compute the distance between matrix and query (FP32, M=1, N=1)
4858
void SquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(const ValueType *m,
4959
const ValueType *q,
5060
size_t dim,
5161
float *out) {
62+
<<<<<<< HEAD
5263
#if defined(__ARM_NEON)
5364
SquaredEuclideanDistanceFp32NEON(m, q, dim, out);
5465
#else
66+
=======
67+
#if defined(__riscv_vector)
68+
if (zvec::ailego::internal::CpuFeatures::static_flags_.RISCV_VECTOR) {
69+
*out = SquaredEuclideanDistanceRVV(m, q, dim);
70+
return;
71+
}
72+
#elif defined(__ARM_NEON)
73+
SquaredEuclideanDistanceNEON(m, q, dim, out);
74+
#elif defined(__SSE__)
75+
>>>>>>> 05b06b8 (feat: add RVV non-batch distance operators)
5576
#if defined(__AVX512F__)
5677
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F) {
5778
*out = SquaredEuclideanDistanceFp32AVX512(m, q, dim);
@@ -64,6 +85,7 @@ void SquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(const ValueType *m,
6485
return;
6586
}
6687
#endif // __AVX__
88+
<<<<<<< HEAD
6789

6890
#if defined(__SSE__)
6991
if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE) {
@@ -74,17 +96,45 @@ void SquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(const ValueType *m,
7496
*out = SquaredEuclideanDistanceFp32Scalar(m, q, dim);
7597
#endif // __ARM_NEON
7698
}
99+
=======
100+
*out = SquaredEuclideanDistanceSSE(m, q, dim);
101+
#else
102+
float sum = 0.0f;
103+
for (size_t i = 0; i < dim; ++i) {
104+
sum += MathHelper::SquaredDifference(m[i], q[i]);
105+
}
106+
*out = sum;
107+
#endif
108+
}
109+
#endif // __SSE__ || __ARM_NEON || __riscv_vector
110+
111+
>>>>>>> 05b06b8 (feat: add RVV non-batch distance operators)
77112

78113
//-----------------------------------------------------------
79114
// EuclideanDistance
80115
//-----------------------------------------------------------
116+
<<<<<<< HEAD
117+
=======
118+
#if defined(__SSE__) || (defined(__ARM_NEON) && defined(__aarch64__)) || \
119+
defined(__riscv_vector)
120+
>>>>>>> 05b06b8 (feat: add RVV non-batch distance operators)
81121
//! Compute the distance between matrix and query (FP32, M=1, N=1)
82122
void EuclideanDistanceMatrix<float, 1, 1>::Compute(const ValueType *m,
83123
const ValueType *q,
84124
size_t dim, float *out) {
125+
#if defined(__riscv_vector)
126+
if (zvec::ailego::internal::CpuFeatures::static_flags_.RISCV_VECTOR) {
127+
*out = EuclideanDistanceRVV(m, q, dim);
128+
return;
129+
}
130+
#endif
85131
SquaredEuclideanDistanceMatrix<float, 1, 1>::Compute(m, q, dim, out);
86132
*out = std::sqrt(*out);
87133
}
134+
<<<<<<< HEAD
135+
=======
136+
#endif // __SSE__ || __ARM_NEON && __aarch64__ || __riscv_vector
137+
>>>>>>> 05b06b8 (feat: add RVV non-batch distance operators)
88138

89139
} // namespace ailego
90-
} // namespace zvec
140+
} // namespace zvec

‎src/ailego/math/euclidean_distance_matrix_fp32_rvv.cc‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Copyright 2025-present the zvec project
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include <cmath>
16+
#include <zvec/ailego/internal/platform.h>
17+
#include "distance_matrix_euclidean_utility.i"
18+
#include "euclidean_distance_matrix.h"
19+
20+
namespace zvec {
21+
namespace ailego {
22+
23+
#if defined(__riscv_vector)
24+
namespace {
25+
26+
static inline float SquaredEuclideanDistanceRVVImpl(const float *lhs,
27+
const float *rhs,
28+
size_t size) {
29+
const size_t vlmax = __riscv_vsetvlmax_e32m8();
30+
vfloat32m8_t v_sum = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
31+
32+
while (size != 0) {
33+
size_t vl = __riscv_vsetvl_e32m8(size);
34+
vfloat32m8_t v_lhs = __riscv_vle32_v_f32m8(lhs, vl);
35+
vfloat32m8_t v_rhs = __riscv_vle32_v_f32m8(rhs, vl);
36+
vfloat32m8_t v_d = __riscv_vfsub_vv_f32m8(v_lhs, v_rhs, vl);
37+
v_sum = __riscv_vfmacc_vv_f32m8_tu(v_sum, v_d, v_d, vl);
38+
lhs += vl;
39+
rhs += vl;
40+
size -= vl;
41+
}
42+
43+
vfloat32m1_t v_zero = __riscv_vfmv_v_f_f32m1(0.0f, 1);
44+
vfloat32m1_t v_red = __riscv_vfredusum_vs_f32m8_f32m1(v_sum, v_zero, vlmax);
45+
return __riscv_vfmv_f_s_f32m1_f32(v_red);
46+
}
47+
48+
} // namespace
49+
50+
//! Squared Euclidean Distance
51+
float SquaredEuclideanDistanceRVV(const float *lhs, const float *rhs,
52+
size_t size) {
53+
return SquaredEuclideanDistanceRVVImpl(lhs, rhs, size);
54+
}
55+
56+
//! Euclidean Distance
57+
float EuclideanDistanceRVV(const float *lhs, const float *rhs, size_t size) {
58+
return std::sqrt(SquaredEuclideanDistanceRVVImpl(lhs, rhs, size));
59+
}
60+
61+
#endif // __riscv_vector
62+
63+
} // namespace ailego
64+
} // namespace zvec

0 commit comments

Comments
 (0)