facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add RVV SIMD specialization for product quantizer distance-table computation (#5543)

Browse Source 
Summary:
Add a RISC-V Vector Extension implementation for the product-quantizer
distance-table computation, and wire it into the SIMD dispatch path. The
kernels are runtime vector-length aware (vsetvl-driven, e32m1) rather than
assuming a fixed vector width.

Changes
- compute_PQ_dis_tables_dsub2<RISCV_RVV> (utils/distances_rvv.cpp): RVV
  kernel for both inner-product and L2 dsub=2 distance tables. Centroids are
  first split into even/odd dimension planes (c0_all / c1_all) so each
  subquantizer's ksub entries load as contiguous vectors; the hot loop
  processes two subquantizers per iteration:
      IP:  r = c0*x0 + c1*x1         (vfmul + vfmacc)
      L2:  r = (c0-x0)^2 + (c1-x1)^2 (vfsub + vfmul + vfadd)
- Add RVV specializations for the accompanying distance primitives:
  fvec_norm_L2sqr, fvec_L2sqr_ny, fvec_L2sqr_ny_transposed and
  fvec_inner_products_ny are vectorized (strided vlse32 over ny); the
  remaining fvec_* and VectorDistance specializations fall back to
  SIMDLevel::NONE.
- compute_distance_tables / compute_inner_prod_tables: route dsub == 2 &&
  nbits >= 3 && nbits < 8 through compute_PQ_dis_tables_dsub2 when
  COMPILE_SIMD_RISCV_RVV is defined, matching the AVX2 / ARM_NEON fast path.
- compute_PQ_dis_tables_dsub2_dispatch: extend the dispatch mask to include
  RISCV_RVV.

### Review fixes

The AI reviewers reported 6 valid defects. This version corrects all of them.

- The RVV `compute_PQ_dis_tables_dsub2` unrolled the subquantizer loop by 2
  and had no tail. An odd `M` left the last `ksub` table entries unwritten.
  The kernel now writes the last subquantizer in a tail loop.
- `fvec_inner_products_ny<RISCV_RVV>` loaded every chunk from `y + j`. The
  base pointer missed the `i * d` offset. Each chunk after the first repeated
  the first vectors. The kernel now loads from `y + i * d + j`.
- `fvec_norm_L2sqr<RISCV_RVV>` ran one `vfredusum` per vector chunk. A
  `d=768` input at VLEN=128 cost 24 serialized reductions. The function again
  keeps a tail-undisturbed accumulator, and it reduces once after the loop.
- The `dsub2` kernels require `ksub % 8 == 0`. The dispatch condition accepted
  every `nbits < 8`, so an `nbits` of 1 or 2 made `compute_distance_tables`
  throw. The condition now also requires `nbits >= 3`. This defect also
  affected the AVX2 and the ARM NEON paths.
- `search_sdc` dropped the per-thread `q_row` pointer precomputation. That
  change added `M` multiplications per database vector on every architecture.
  This version restores the precomputation.
- The `ny` kernels fixed the chunk width at the VLEN=128 value. They now read
  `VLMAX` with `vsetvlmax`. Hardware with a larger VLEN uses the whole vector
  register.

The diff also removes 2 unreadable comments from the RVV kernel.

Caveat: the benchmark below ran on the previous version. That version included
the `search_sdc` change that this version reverts.

### Performance (sg2044)
Cmd: 1M 768 <PQ> 32 100 16 1 1000 1 <metric>  (search ns/query, lower is better)

| #  | CONFIG   | M   | NPROBE | METRIC | F1(ns)  | PQ1(ns) | F2(ns)  | PQ2(ns) | Δ_PQ%   | Δ_F%   |
|----|----------|-----|--------|--------|---------|---------|---------|---------|---------|--------|
| 01 | PQ8x4    | 4   | 8      | L2     | 6656870 | 536988  | 6582134 | 495970  | -7.64%  | -1.12% |
| 02 | PQ8x4    | 4   | 8      | IP     | 6228497 | 475320  | 6164081 | 481118  | +1.22%  | -1.03% |
| 03 | PQ8x8    | 8   | 8      | L2     | 6626171 | 545230  | 6580586 | 500343  | -8.23%  | -0.69% |
| 04 | PQ8x8    | 8   | 8      | IP     | 6228452 | 519710  | 6165707 | 479864  | -7.67%  | -1.01% |
| 05 | PQ8x16   | 16  | 8      | L2     | 6615305 | 697070  | 6582239 | 624059  | -10.47% | -0.50% |
| 06 | PQ8x16   | 16  | 8      | IP     | 6227067 | 642949  | 6164572 | 587271  | -8.66%  | -1.00% |
| 07 | PQ8x32   | 32  | 8      | L2     | 6624323 | 1185655 | 6581329 | 902298  | -23.90% | -0.65% |
| 08 | PQ8x32   | 32  | 8      | IP     | 6239930 | 1155545 | 6164198 | 785403  | -32.03% | -1.21% |
| 09 | PQ8x64   | 64  | 8      | L2     | 6624196 | 2006275 | 6578815 | 1687191 | -15.90% | -0.69% |
| 10 | PQ8x64   | 64  | 8      | IP     | 6225719 | 1800678 | 6164767 | 1512760 | -15.99% | -0.98% |
| 11 | PQ8x96   | 96  | 8      | L2     | 6617525 | 3108377 | 6582946 | 2701758 | -13.08% | -0.52% |
| 12 | PQ8x96   | 96  | 8      | IP     | 6239711 | 2657958 | 6170006 | 2320803 | -12.68% | -1.12% |
| 13 | PQ8x384  | 384 | 8      | L2     | 6620602 | 14029275| 6583820 | 12318712| -12.19% | -0.56% |
| 14 | PQ8x384  | 384 | 8      | IP     | 6228776 | 12089134| 6166143 | 10509785| -13.06% | -1.01% |
| 15 | PQ8x32   | 32  | 4      | L2     | 3321625 | 773663  | 3298874 | 491225  | -36.51% | -0.68% |
| 16 | PQ8x32   | 32  | 16     | L2     |13223858 | 2184906 |13152640 | 1645444 | -24.69% | -0.54% |
| 17 | PQ8x32   | 32  | 32     | L2     |26430233 | 3405062 |26331426 | 3138563 | -7.83%  | -0.37% |
| 18 | PQ8x64   | 64  | 4      | L2     | 3320850 | 1191054 | 3298426 | 969309  | -18.62% | -0.68% |
| 19 | PQ8x64   | 64  | 16     | L2     |13207340 | 3944868 |13149902 | 3252654 | -17.55% | -0.43% |
| 20 | PQ8x64   | 64  | 32     | L2     |26467696 | 6725098 |26335327 | 6667493 | -0.86%  | -0.50% |
| 21 | PQ8x16   | 16  | 4      | L2     | 3316496 | 427653  | 3304051 | 367407  | -14.09% | -0.38% |
| 22 | PQ8x16   | 16  | 16     | L2     |13203491 | 1198002 |13150061 | 1121421 | -6.39%  | -0.40% |
| 23 | PQ8x16   | 16  | 32     | L2     |26438401 | 2199527 |26280149 | 2143973 | -2.53%  | -0.60% |

Δ_PQ%: PQ baseline -> RVV; Δ_F%: Flat baseline -> RVV.

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

Test Plan:
`buck build fbcode//faiss:faiss` passes.

`buck2 test` passes on the PQ and the distance C++ tests:
`test_disable_pq_sdc_tables`, `test_distances_dispatch`, `test_distances_simd`,
`test_ivfpq_codec`, `test_ivfpq_indexing`, `test_pq_encoding` and
`test_pq_code_distance`. The result is 27 pass and 0 fail. The host is aarch64,
so 8 tests skip. Those tests need a second SIMD level to compare against.

This diff adds 2 tests to `faiss/tests/test_product_quantizer.py`:
- `test_dsub2_4bit` covers `dsub=2` with `nbits=4`. `ksub` is 16, so the SIMD
  `dsub2` kernel runs.
- `test_dsub2_2bit` covers `dsub=2` with `nbits=2`. `ksub` is 4, so the generic
  path must run instead.

I checked the `nbits >= 3` fix with a temporary C++ test on ARM NEON. The test
compares `compute_distance_tables` and `compute_inner_prod_tables` against
`compute_distance_table` and `compute_inner_prod_table`, for `nbits` of 1, 2, 3
and 4. The test fails with a `ksub % 8 == 0` assertion before the fix. The test
passes after the fix. I removed the temporary test before submit.

Limits of the local verification:
- I did not run `faiss/tests/test_product_quantizer.py` on this devserver. The
  `fbcode//faiss/python:pyfaiss-py-gen` SWIG step fails on aarch64 with
  `Exec format error`. The failure also reproduces without my changes. CI runs
  the Python tests on x86.
- I did not compile `utils/simd_impl/distances_rvv.cpp`. This devserver has no
  RISC-V toolchain. `xplat.bzl` registers no RVV sources, so Buck and internal
  CI do not build that file either. Only the OSS CMake build on rv64 compiles
  it. A reviewer with RVV hardware should confirm the 4 RVV kernel fixes.

Reviewed By: juancarpio27

Differential Revision: D117914440

Pulled By: mnorris11

fbshipit-source-id: 3539a116d47f4a96c1051bc10e6440deae8c5302
This commit is contained in:
anthony-zy
2026-09-24 18:31:50 -07:00
committed by meta-codesync[bot]
parent b1c1496545
commit 1768a5a8e9
4 changed files with 271 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/impl/ProductQuantizer.cpp
+13 -6
View File
@@ -489,8 +489,12 @@ void ProductQuantizer::compute_distance_tables(
const float* x,
float* dis_tables) const {
int64_t nx_signed = nx;
#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON)
if (dsub == 2 && nbits < 8) { // interesting for a narrow range of settings
#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON) || \
defined(COMPILE_SIMD_RISCV_RVV)
// The dsub2 kernels help for a narrow range of settings. They require
// ksub = 1 << nbits to be a multiple of 8. Therefore nbits must be 3 or
// more.
if (dsub == 2 && nbits >= 3 && nbits < 8) {
compute_PQ_dis_tables_dsub2(
d, ksub, centroids.data(), nx, x, false, dis_tables);
} else
@@ -524,8 +528,9 @@ void ProductQuantizer::compute_inner_prod_tables(
const float* x,
float* dis_tables) const {
int64_t nx_signed = nx;
#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON)
if (dsub == 2 && nbits < 8) {
#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON) || \
defined(COMPILE_SIMD_RISCV_RVV)
if (dsub == 2 && nbits >= 3 && nbits < 8) {
compute_PQ_dis_tables_dsub2(
d, ksub, centroids.data(), nx, x, true, dis_tables);
} else
@@ -872,8 +877,9 @@ void ProductQuantizer::search_sdc(
float* heap_dis = res->val + i * k;
const uint8_t* qcode = qcodes + i * code_size;
if (init_finalize_heap)
if (init_finalize_heap) {
maxheap_heapify(k, heap_dis, heap_ids);
}
// Precompute per-subquantizer row pointers: q_row[m] points to
// sdc_table[m*ksub^2 + qcode[m]*ksub], eliminating M
@@ -897,8 +903,9 @@ void ProductQuantizer::search_sdc(
bcode += code_size;
}
if (init_finalize_heap)
if (init_finalize_heap) {
maxheap_reorder(k, heap_dis, heap_ids);
}
}
}
}
faiss/utils/distances_dispatch.h
+3 -1
View File
@@ -214,7 +214,9 @@ inline void compute_PQ_dis_tables_dsub2_dispatch(
const float* x,
bool is_inner_product,
float* dis_tables) {
with_simd_level_256bit([&]<SIMDLevel level>() {
constexpr int dsub2_mask =
AVAILABLE_SIMD_LEVELS_AVX2_NEON | (1 << int(SIMDLevel::RISCV_RVV));
with_selected_simd_levels<dsub2_mask>([&]<SIMDLevel level>() {
compute_PQ_dis_tables_dsub2<level>(
d, ksub, centroids, nx, x, is_inner_product, dis_tables);
});
faiss/utils/simd_impl/distances_rvv.cpp
+240 -34
View File
@@ -13,6 +13,7 @@
#include <faiss/utils/extra_distances.h>
#include <riscv_vector.h>
#include <vector>
namespace faiss {
@@ -52,6 +53,126 @@ static inline size_t rvv_argmin(const float* values, size_t n) {
return n;
}
// Write the ksub table entries of one subquantizer that has dsub = 2.
// The function computes one of these 2 results:
// inner product: out[k] = x0 * c0[k] + x1 * c1[k]
// L2: out[k] = (c0[k] - x0)^2 + (c1[k] - x1)^2
template <bool is_inner_product>
static inline void rvv_pq2_row(
const float* c0,
const float* c1,
float x0,
float x1,
size_t ksub,
size_t vl,
float* out) {
for (size_t k = 0; k < ksub; k += vl) {
vfloat32m1_t vc0 = __riscv_vle32_v_f32m1(c0 + k, vl);
vfloat32m1_t vc1 = __riscv_vle32_v_f32m1(c1 + k, vl);
if constexpr (is_inner_product) {
vfloat32m1_t r = __riscv_vfmacc_vf_f32m1(
__riscv_vfmul_vf_f32m1(vc0, x0, vl), x1, vc1, vl);
__riscv_vse32_v_f32m1(out + k, r, vl);
} else {
vfloat32m1_t d0 = __riscv_vfsub_vf_f32m1(vc0, x0, vl);
vfloat32m1_t d1 = __riscv_vfsub_vf_f32m1(vc1, x1, vl);
d0 = __riscv_vfmul_vv_f32m1(d0, d0, vl);
d1 = __riscv_vfmul_vv_f32m1(d1, d1, vl);
vfloat32m1_t r = __riscv_vfadd_vv_f32m1(d0, d1, vl);
__riscv_vse32_v_f32m1(out + k, r, vl);
}
}
}
template <bool is_inner_product>
static void rvv_pq2_tables(
size_t M,
size_t ksub,
const float* c0_all,
const float* c1_all,
size_t nx,
const float* x,
size_t vl,
float* dis_tables) {
const size_t d = 2 * M;
for (size_t i = 0; i < nx; i++) {
const float* xi = x + i * d;
float* oi = dis_tables + i * M * ksub;
size_t m = 0;
// The loop computes 2 subquantizers per iteration. The 2 computations
// are independent, so the processor can run them at the same time.
for (; m + 1 < M; m += 2) {
rvv_pq2_row<is_inner_product>(
c0_all + m * ksub,
c1_all + m * ksub,
xi[2 * m],
xi[2 * m + 1],
ksub,
vl,
oi + m * ksub);
rvv_pq2_row<is_inner_product>(
c0_all + (m + 1) * ksub,
c1_all + (m + 1) * ksub,
xi[2 * (m + 1)],
xi[2 * (m + 1) + 1],
ksub,
vl,
oi + (m + 1) * ksub);
}
// If M is odd, the loop above skips the last subquantizer. This loop
// computes it.
for (; m < M; m++) {
rvv_pq2_row<is_inner_product>(
c0_all + m * ksub,
c1_all + m * ksub,
xi[2 * m],
xi[2 * m + 1],
ksub,
vl,
oi + m * ksub);
}
}
}
template <>
void compute_PQ_dis_tables_dsub2<SIMDLevel::RISCV_RVV>(
size_t d,
size_t ksub,
const float* all_centroids,
size_t nx,
const float* x,
bool is_inner_product,
float* dis_tables) {
const size_t M = d / 2;
FAISS_THROW_IF_NOT(ksub % 8 == 0);
// Separate the centroids into 2 planes. Each plane holds one component of
// every centroid, so a unit-stride load can read the plane.
std::vector<float> c0_all(M * ksub);
std::vector<float> c1_all(M * ksub);
for (size_t m = 0; m < M; m++) {
const float* cm = all_centroids + m * ksub * 2;
float* c0m = c0_all.data() + m * ksub;
float* c1m = c1_all.data() + m * ksub;
for (size_t k = 0; k < ksub; k++) {
c0m[k] = cm[2 * k];
c1m[k] = cm[2 * k + 1];
}
}
// ksub is a power of 2, and ksub is a multiple of 8. Therefore vl divides
// ksub exactly.
const size_t vl = __riscv_vsetvl_e32m1(ksub);
if (is_inner_product) {
rvv_pq2_tables<true>(
M, ksub, c0_all.data(), c1_all.data(), nx, x, vl, dis_tables);
} else {
rvv_pq2_tables<false>(
M, ksub, c0_all.data(), c1_all.data(), nx, x, vl, dis_tables);
}
}
template <>
float fvec_norm_L2sqr<SIMDLevel::RISCV_RVV>(const float* x, size_t d) {
size_t vlmax = __riscv_vsetvlmax_e32m8();
@@ -227,31 +348,58 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
size_t d,
size_t d_offset,
size_t 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;
// Compute the squared length of the query subvector.
float x_sqlen = 0;
for (size_t j = 0; j < d; j++) {
x_sqlen += x[j] * x[j];
}
float x_sqlen =
rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
i = 0;
while (i < ny) {
size_t vl = __riscv_vsetvl_e32m8(ny - i);
acc = __riscv_vfmv_v_f_f32m8(0.0f, vl);
// The loop below divides the ny dimension into chunks. Each chunk uses
// e32m8. The chunk size is VLMAX, so the kernel fills the whole vector
// register at every VLEN.
const size_t chunk = __riscv_vsetvlmax_e32m8();
size_t i = 0;
// dis[i] = x_sqlen + y_sqlen[i] - 2 * sum_j x[j] * y[j * d_offset + i]
//
// The code accumulates the dot product first, and it applies the factor
// of 2 afterwards. A factor of -2 on x[j] can overflow to an infinity.
// That infinity then multiplies an element of y that is 0, and the
// product is a NaN. The SIMDLevel::NONE implementation scales the
// accumulated dot product in the same way.
for (; i + chunk <= ny; i += chunk) {
vfloat32m8_t dp = __riscv_vfmv_v_f_f32m8(0.0f, chunk);
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 y_vec =
__riscv_vle32_v_f32m8(y + j * d_offset + i, chunk);
dp = __riscv_vfmacc_vf_f32m8(dp, x[j], y_vec, chunk);
}
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;
vfloat32m8_t acc = __riscv_vle32_v_f32m8(y_sqlen + i, chunk);
acc = __riscv_vfadd_vf_f32m8(acc, x_sqlen, chunk);
dp = __riscv_vfmul_vf_f32m8(dp, 2.0f, chunk);
acc = __riscv_vfsub_vv_f32m8(acc, dp, chunk);
__riscv_vse32_v_f32m8(dis + i, acc, chunk);
}
// This block computes the last ny % chunk elements.
if (i < ny) {
size_t vl = __riscv_vsetvl_e32m8(ny - i);
vfloat32m8_t dp = __riscv_vfmv_v_f_f32m8(0.0f, vl);
for (size_t j = 0; j < d; j++) {
vfloat32m8_t y_vec =
__riscv_vle32_v_f32m8(y + j * d_offset + i, vl);
dp = __riscv_vfmacc_vf_f32m8(dp, x[j], y_vec, vl);
}
vfloat32m8_t acc = __riscv_vle32_v_f32m8(y_sqlen + i, vl);
acc = __riscv_vfadd_vf_f32m8(acc, x_sqlen, vl);
dp = __riscv_vfmul_vf_f32m8(dp, 2.0f, vl);
acc = __riscv_vfsub_vv_f32m8(acc, dp, vl);
__riscv_vse32_v_f32m8(dis + i, acc, vl);
}
}
@@ -262,9 +410,43 @@ void fvec_inner_products_ny<SIMDLevel::RISCV_RVV>(
const float* y,
size_t d,
size_t ny) {
for (size_t i = 0; i < ny; i++) {
ip[i] = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y, d);
y += d;
// The loop below divides the ny dimension into chunks. Each chunk uses
// e32m4. An LMUL sweep showed that e32m4 is the fastest choice. It
// balances the throughput of the strided load vlse32 against the use of
// the register file. That load uses a stride of d*4 bytes. The chunk size
// is VLMAX, so the kernel fills the whole vector register at every VLEN.
// ip[i] = sum_{j=0}^{d-1} x[j] * y[i * d + j]
const size_t chunk = __riscv_vsetvlmax_e32m4();
const ptrdiff_t stride_bytes = (ptrdiff_t)d * (ptrdiff_t)sizeof(float);
size_t i = 0;
for (; i + chunk <= ny; i += chunk) {
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, chunk);
const float* yb = y + i * d; // the first vector of this chunk
for (size_t j = 0; j < d; j++) {
vfloat32m4_t y_vec =
__riscv_vlse32_v_f32m4(yb + j, stride_bytes, chunk);
acc = __riscv_vfmacc_vf_f32m4(acc, x[j], y_vec, chunk);
}
__riscv_vse32_v_f32m4(ip + i, acc, chunk);
}
// This block computes the last ny % chunk elements.
if (i < ny) {
size_t vl = __riscv_vsetvl_e32m4(ny - i);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
const float* yb = y + i * d;
for (size_t j = 0; j < d; j++) {
vfloat32m4_t y_vec =
__riscv_vlse32_v_f32m4(yb + j, stride_bytes, vl);
acc = __riscv_vfmacc_vf_f32m4(acc, x[j], y_vec, vl);
}
__riscv_vse32_v_f32m4(ip + i, acc, vl);
}
}
@@ -275,9 +457,36 @@ void fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(
const float* y,
size_t d,
size_t ny) {
for (size_t i = 0; i < ny; i++) {
dis[i] = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y, d);
y += d;
// The loop below divides the ny dimension into chunks. Each chunk uses
// e32m8. The chunk size is VLMAX, so the kernel fills the whole vector
// register at every VLEN.
const size_t chunk = __riscv_vsetvlmax_e32m8();
const ptrdiff_t stride_bytes = (ptrdiff_t)d * (ptrdiff_t)sizeof(float);
size_t i = 0;
for (; i + chunk <= ny; i += chunk) {
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, chunk);
const float* yb = y + i * d; // the first vector of this chunk
for (size_t j = 0; j < d; j++) {
vfloat32m8_t y_vec =
__riscv_vlse32_v_f32m8(yb + j, stride_bytes, chunk);
vfloat32m8_t diff = __riscv_vfsub_vf_f32m8(y_vec, x[j], chunk);
acc = __riscv_vfmacc_vv_f32m8(acc, diff, diff, chunk);
}
__riscv_vse32_v_f32m8(dis + i, acc, chunk);
}
if (i < ny) {
size_t vl = __riscv_vsetvl_e32m8(ny - i);
vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vl);
const float* yb = y + i * d;
for (size_t j = 0; j < d; j++) {
vfloat32m8_t y_vec =
__riscv_vlse32_v_f32m8(yb + j, stride_bytes, vl);
vfloat32m8_t diff = __riscv_vfsub_vf_f32m8(y_vec, x[j], vl);
acc = __riscv_vfmacc_vv_f32m8(acc, diff, diff, vl);
}
__riscv_vse32_v_f32m8(dis + i, acc, vl);
}
}
@@ -288,9 +497,8 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::RISCV_RVV>(
const float* y,
size_t d,
size_t 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;
return fvec_L2sqr_ny_nearest<SIMDLevel::NONE>(
distances_tmp_buffer, x, y, d, ny);
}
template <>
@@ -302,10 +510,8 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::RISCV_RVV>(
size_t d,
size_t d_offset,
size_t ny) {
fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
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 <>
tests/test_product_quantizer.py
+15
View File
@@ -212,3 +212,18 @@ class TestPQTables(unittest.TestCase):
def test_4bit(self):
self.do_test(32, 4, nbit=4)
def test_dsub2_4bit(self):
# ksub = 16, so the SIMD dsub2 kernels run. M = 8 is even.
self.do_test(16, 2, nbit=4)
def test_dsub2_4bit_odd(self):
# M = 9 is odd, so the dsub2 kernels must also write the last
# subquantizer. test_dsub2_odd cannot cover this case, because
# nbit=8 does not reach the dsub2 kernels.
self.do_test(18, 2, nbit=4)
def test_dsub2_2bit(self):
# ksub = 4 is not a multiple of 8, so the dsub2 kernels must be
# skipped in favour of the generic path
self.do_test(16, 2, nbit=2)