facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(riscv): Add RISC-V RVV vector distance kernels (#5354)

Browse Source 
Summary:
This PR adds RISC-V RVV implementations for the float vector distance utilities under `COMPILE_SIMD_RISCV_RVV`.

Implemented RVV specializations include:

* `fvec_norm_L2sqr`
* `fvec_L2sqr`
* `fvec_inner_product`
* `fvec_L1`
* `fvec_Linf`
* `fvec_inner_product_batch_4`
* `fvec_L2sqr_batch_4`
* `fvec_L2sqr_ny_transposed`
* `fvec_inner_products_ny`
* `fvec_L2sqr_ny`
* `fvec_L2sqr_ny_nearest`
* `fvec_L2sqr_ny_nearest_y_transposed`
* `fvec_madd`
* `fvec_madd_and_argmin`

## Testing

Tested on SG2044 with:

* GCC 15.1
* openEuler
* RISC-V RVV enabled build

The RVV results were validated against the scalar implementation, and numerical accuracy is aligned with the scalar path.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5354

Reviewed By: mnorris11

Differential Revision: D110116107

Pulled By: limqiying

fbshipit-source-id: cf6d8e205cfc5f6679c21f1a4e41c0f310d4a315
This commit is contained in:
ihb2032
2026-06-30 14:45:55 -07:00
committed by meta-codesync[bot]
parent aa9d2bf425
commit df51012a60
1 changed files with 160 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/utils/simd_impl/distances_rvv.cpp
+160 -18
View File
@@ -12,12 +12,24 @@
#ifdef COMPILE_SIMD_RISCV_RVV
#include <faiss/utils/extra_distances.h>
#include <riscv_vector.h>
namespace faiss {
template <>
float fvec_norm_L2sqr<SIMDLevel::RISCV_RVV>(const float* x, size_t d) {
return fvec_norm_L2sqr<SIMDLevel::NONE>(x, d);
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
size_t i = 0;
while (i < d) {
size_t vl = __riscv_vsetvl_e32m8(d - i);
vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
i += vl;
}
vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(sum);
}
template <>
@@ -25,7 +37,20 @@ float fvec_L2sqr<SIMDLevel::RISCV_RVV>(
const float* x,
const float* y,
size_t d) {
return fvec_L2sqr<SIMDLevel::NONE>(x, y, d);
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
size_t i = 0;
while (i < d) {
size_t vl = __riscv_vsetvl_e32m8(d - i);
vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
i += vl;
}
vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(sum);
}
template <>
@@ -33,12 +58,38 @@ float fvec_inner_product<SIMDLevel::RISCV_RVV>(
const float* x,
const float* y,
size_t d) {
return fvec_inner_product<SIMDLevel::NONE>(x, y, d);
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
size_t i = 0;
while (i < d) {
size_t vl = __riscv_vsetvl_e32m8(d - i);
vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vy, vl);
i += vl;
}
vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(sum);
}
template <>
float fvec_L1<SIMDLevel::RISCV_RVV>(const float* x, const float* y, size_t d) {
return fvec_L1<SIMDLevel::NONE>(x, y, d);
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
size_t i = 0;
while (i < d) {
size_t vl = __riscv_vsetvl_e32m8(d - i);
vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
vx = __riscv_vfsgnjx_vv_f32m8(vx, vx, vl);
acc = __riscv_vfadd_vv_f32m8_tu(acc, acc, vx, vl);
i += vl;
}
vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(sum);
}
template <>
@@ -46,7 +97,21 @@ float fvec_Linf<SIMDLevel::RISCV_RVV>(
const float* x,
const float* y,
size_t d) {
return fvec_Linf<SIMDLevel::NONE>(x, y, d);
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t vmax = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
size_t i = 0;
while (i < d) {
size_t vl = __riscv_vsetvl_e32m8(d - i);
vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
vx = __riscv_vfsgnjx_vv_f32m8(vx, vx, vl);
vmax = __riscv_vfmax_vv_f32m8_tu(vmax, vmax, vx, vl);
i += vl;
}
vfloat32m1_t max = __riscv_vfmv_s_f_f32m1(0.0f, 1);
max = __riscv_vfredmax_vs_f32m8_f32m1(vmax, max, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(max);
}
template <>
@@ -61,8 +126,10 @@ void fvec_inner_product_batch_4<SIMDLevel::RISCV_RVV>(
float& dis1,
float& dis2,
float& dis3) {
fvec_inner_product_batch_4<SIMDLevel::NONE>(
x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
dis0 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y0, d);
dis1 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y1, d);
dis2 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y2, d);
dis3 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y3, d);
}
template <>
@@ -77,8 +144,10 @@ void fvec_L2sqr_batch_4<SIMDLevel::RISCV_RVV>(
float& dis1,
float& dis2,
float& dis3) {
fvec_L2sqr_batch_4<SIMDLevel::NONE>(
x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
dis0 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y0, d);
dis1 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y1, d);
dis2 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y2, d);
dis3 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y3, d);
}
template <>
@@ -90,8 +159,33 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
size_t d,
size_t d_offset,
size_t ny) {
fvec_L2sqr_ny_transposed<SIMDLevel::NONE>(
dis, x, y, y_sqlen, d, d_offset, ny);
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
size_t i = 0;
while (i < d) {
size_t vl = __riscv_vsetvl_e32m8(d - i);
vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
i += vl;
}
vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
float x_sqlen = __riscv_vfmv_f_s_f32m1_f32(sum);
i = 0;
while (i < ny) {
size_t vl = __riscv_vsetvl_e32m8(ny - i);
acc = __riscv_vfmv_v_f_f32m8(0.0f, vl);
for (size_t j = 0; j < d; j++) {
vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + j * d_offset + i, vl);
acc = __riscv_vfmacc_vf_f32m8(acc, x[j], vy, vl);
}
vfloat32m8_t vres = __riscv_vle32_v_f32m8(y_sqlen + i, vl);
vres = __riscv_vfadd_vf_f32m8(vres, x_sqlen, vl);
acc = __riscv_vfmul_vf_f32m8(acc, 2.0f, vl);
vres = __riscv_vfsub_vv_f32m8(vres, acc, vl);
__riscv_vse32_v_f32m8(dis + i, vres, vl);
i += vl;
}
}
template <>
@@ -101,7 +195,10 @@ void fvec_inner_products_ny<SIMDLevel::RISCV_RVV>(
const float* y,
size_t d,
size_t ny) {
fvec_inner_products_ny<SIMDLevel::NONE>(ip, x, y, d, ny);
for (size_t i = 0; i < ny; i++) {
ip[i] = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y, d);
y += d;
}
}
template <>
@@ -111,7 +208,39 @@ void fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(
const float* y,
size_t d,
size_t ny) {
fvec_L2sqr_ny<SIMDLevel::NONE>(dis, x, y, d, ny);
for (size_t i = 0; i < ny; i++) {
dis[i] = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y, d);
y += d;
}
}
// Index of the first element equal to the minimum of values[0..n), or n when
// there is none (e.g. n == 0). Shared by the *_nearest and madd_and_argmin
// kernels so the vfmin/vfredmin/vfirst sequence lives in one place.
static size_t rvv_argmin(const float* values, size_t n) {
size_t vlmax = __riscv_vsetvlmax_e32m8();
vfloat32m8_t vmin = __riscv_vfmv_v_f_f32m8(__builtin_inff(), vlmax);
size_t i = 0;
while (i < n) {
size_t vl = __riscv_vsetvl_e32m8(n - i);
vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
vmin = __riscv_vfmin_vv_f32m8_tu(vmin, vmin, vd, vl);
i += vl;
}
vfloat32m1_t rmin = __riscv_vfmv_s_f_f32m1(__builtin_inff(), 1);
rmin = __riscv_vfredmin_vs_f32m8_f32m1(vmin, rmin, vlmax);
float min_val = __riscv_vfmv_f_s_f32m1_f32(rmin);
i = 0;
while (i < n) {
size_t vl = __riscv_vsetvl_e32m8(n - i);
vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
long j = __riscv_vfirst_m_b4(
__riscv_vmfeq_vf_f32m8_b4(vd, min_val, vl), vl);
if (j >= 0)
return i + static_cast<size_t>(j);
i += vl;
}
return n;
}
template <>
@@ -121,8 +250,9 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::RISCV_RVV>(
const float* y,
size_t d,
size_t ny) {
return fvec_L2sqr_ny_nearest<SIMDLevel::NONE>(
distances_tmp_buffer, x, y, d, ny);
fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(distances_tmp_buffer, x, y, d, ny);
const size_t j = rvv_argmin(distances_tmp_buffer, ny);
return j < ny ? j : 0;
}
template <>
@@ -134,8 +264,10 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::RISCV_RVV>(
size_t d,
size_t d_offset,
size_t ny) {
return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
const size_t j = rvv_argmin(distances_tmp_buffer, ny);
return j < ny ? j : 0;
}
template <>
@@ -145,7 +277,15 @@ void fvec_madd<SIMDLevel::RISCV_RVV>(
float bf,
const float* b,
float* c) {
fvec_madd<SIMDLevel::NONE>(n, a, bf, b, c);
size_t i = 0;
while (i < n) {
size_t vl = __riscv_vsetvl_e32m8(n - i);
vfloat32m8_t va = __riscv_vle32_v_f32m8(a + i, vl);
vfloat32m8_t vb = __riscv_vle32_v_f32m8(b + i, vl);
va = __riscv_vfmacc_vf_f32m8(va, bf, vb, vl);
__riscv_vse32_v_f32m8(c + i, va, vl);
i += vl;
}
}
template <>
@@ -155,7 +295,9 @@ int fvec_madd_and_argmin<SIMDLevel::RISCV_RVV>(
float bf,
const float* b,
float* c) {
return fvec_madd_and_argmin<SIMDLevel::NONE>(n, a, bf, b, c);
fvec_madd<SIMDLevel::RISCV_RVV>(n, a, bf, b, c);
const size_t j = rvv_argmin(c, n);
return j < n ? static_cast<int>(j) : -1;
}
#define DEFINE_VECTOR_DISTANCE_RVV_FALLBACK(metric) \