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

Move the RVV distance kernels into free-standing headers (#5697)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5697

This diff moves the leaf RISC-V Vector (RVV) distance kernels into three free-standing headers. A caller that cannot include the full Faiss headers can then run the same arithmetic. The diff also makes `utils/simd_impl/distances_rvv.cpp` compile with clang, and it removes a duplicate gather loop. The CPU paths do not change behavior.

## The headers

| header | caller |
| --- | --- |
| `impl/scalar_quantizer/sq8_rvv_kernel.h` | `sq-rvv.cpp` |
| `impl/pq_code_distance/pq_rvv_kernel.h` | `pq_code_distance/rvv.cpp` |
| `utils/simd_impl/distances_rvv_kernel.h` | `utils/simd_impl/distances_rvv.cpp` |

Each header holds only the innermost distance loops. It needs no exceptions, no heap and no threads, so a freestanding RISC-V target can compile it. `impl/scalar_quantizer/sq8_coefficients.h` holds the QT_8bit and QT_4bit scale and offset, which any architecture can use. The scalar quantizer header also holds block, two-block, strided and 4-bit forms of the loop, for callers that store codes transposed in blocks of 64.

## The clang fix

Clang gave 14 instances of `error: builtin functions must be directly called` for `utils/simd_impl/distances_rvv.cpp`, because `rvv_reduce` takes the reduction intrinsic as a template argument. `FAISS_RVV_REDUCE_OP` now wraps the intrinsic in a lambda. GCC still accepts the code.

## The duplicate

`pq_code_distance_8bit_single_impl` held its own copy of the gather loop. It now calls `pq_rvv_kernel::distance_8bit`. The four-code entry point keeps its own loop, because it computes the row term once for four codes.

`xplat.bzl` and `CMakeLists.txt` register the four new headers.

Differential Revision: D122938441

fbshipit-source-id: 88e59f674a82d25a60abc18ab330904d6716566b
This commit is contained in:
Michael Norris
2026-10-05 12:43:09 -07:00
committed by meta-codesync[bot]
parent d380b7d96c
commit 83ae8b0908
8 changed files with 876 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+4
View File
@@ -333,6 +333,10 @@ set(FAISS_HEADERS
impl/scalar_quantizer/quantizers.h
impl/scalar_quantizer/scanners.h
impl/scalar_quantizer/similarities.h
impl/pq_code_distance/pq_rvv_kernel.h
impl/scalar_quantizer/sq8_coefficients.h
impl/scalar_quantizer/sq8_rvv_kernel.h
utils/simd_impl/distances_rvv_kernel.h
impl/scalar_quantizer/sq-dispatch.h
impl/scalar_quantizer/training.h
impl/ThreadedIndex-inl.h
faiss/impl/pq_code_distance/pq_rvv_kernel.h
+89
View File
@@ -0,0 +1,89 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
/* 8-bit product-quantizer code distance for RISC-V Vector.
*
* Free-standing (no exceptions, heap or threads), so a freestanding RISC-V
* target can compile it. pq_code_distance/rvv.cpp calls it; no Buck target
* here builds that RISC-V CPU path.
*
* PQDecoder8 fixes nbits = 8, so ksub = 256 and the scalar loop is
*
* for (m = 0; m < M; m++) result += sim_table[m * 256 + code[m]];
*
* That is a gather, not a contiguous read. RVV issues one vluxei32 per strip:
* build the element offset m * 256 + code[m] across the lanes, scale it to
* bytes, then index sim_table.
*
* The offset stays inside uint32 while M * 256 * 4 bytes does. At M = 1024
* that is 1 MiB, so no PQ shape in Faiss overflows it.
*
* The accumulator is LMUL 4, as in utils/simd_impl/distances_rvv_kernel.h.
* The lanes add the M terms in a different order than the scalar loop, so
* the last bits can differ.
*/
#include <cstddef>
#include <cstdint>
#include <riscv_vector.h>
namespace faiss {
namespace pq_rvv_kernel {
/// Element offsets into sim_table for subquantizers [m, m + vl).
inline vuint32m4_t row_offsets(size_t m, size_t vl) {
// (m + lane) * 256
vuint32m4_t row = __riscv_vadd_vx_u32m4(__riscv_vid_v_u32m4(vl), m, vl);
return __riscv_vsll_vx_u32m4(row, 8, vl);
}
/// Byte offsets for one code, given the precomputed row offsets.
inline vuint32m4_t byte_offsets(
vuint32m4_t row_off,
const uint8_t* code,
size_t m,
size_t vl) {
vuint8m1_t cb = __riscv_vle8_v_u8m1(code + m, vl);
vuint32m4_t idx = __riscv_vzext_vf4_u32m4(cb, vl);
// (row + code) * sizeof(float)
return __riscv_vsll_vx_u32m4(
__riscv_vadd_vv_u32m4(row_off, idx, vl), 2, vl);
}
inline float reduce(vfloat32m4_t acc, size_t vlmax) {
vfloat32m1_t zero = __riscv_vfmv_v_f_f32m1(0.0f, 1);
vfloat32m1_t sum = __riscv_vfredusum_vs_f32m4_f32m1(acc, zero, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(sum);
}
/// Distance between a query and one 8-bit PQ code.
/// @param sim_table M * 256 floats, laid out (M, ksub)
inline float distance_8bit(
size_t M,
const float* sim_table,
const uint8_t* code) {
const size_t vlmax = __riscv_vsetvlmax_e32m4();
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
size_t m = 0;
while (m < M) {
const size_t vl = __riscv_vsetvl_e32m4(M - m);
vuint32m4_t off = byte_offsets(row_offsets(m, vl), code, m, vl);
vfloat32m4_t v = __riscv_vluxei32_v_f32m4(sim_table, off, vl);
// Tail-undisturbed: lanes past vl keep earlier partial sums, and the
// final reduction covers vlmax.
acc = __riscv_vfadd_vv_f32m4_tu(acc, acc, v, vl);
m += vl;
}
return reduce(acc, vlmax);
}
} // namespace pq_rvv_kernel
} // namespace faiss
faiss/impl/pq_code_distance/rvv.cpp
+46 -15
View File
@@ -7,12 +7,15 @@
#ifdef COMPILE_SIMD_RISCV_RVV
#include <riscv_vector.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
#include <faiss/impl/pq_code_distance/pq_rvv_kernel.h>
namespace faiss {
namespace pq_code_distance {
// RISCV_RVV: no RVV-optimized PQ code distance exists yet. Use scalar.
// The kernel lives in pq_rvv_kernel.h. See that header for the gather layout.
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
@@ -20,8 +23,7 @@ float pq_code_distance_8bit_single_impl<SIMDLevel::RISCV_RVV>(
size_t M,
const float* sim_table,
const uint8_t* code) {
return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
M, 8, sim_table, code);
return pq_rvv_kernel::distance_8bit(M, sim_table, code);
}
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
@@ -37,18 +39,47 @@ void pq_code_distance_8bit_four_impl<SIMDLevel::RISCV_RVV>(
float& result1,
float& result2,
float& result3) {
PQCodeDistanceScalar<PQDecoder8>::distance_four_codes(
M,
8,
sim_table,
code0,
code1,
code2,
code3,
result0,
result1,
result2,
result3);
const size_t vlmax = __riscv_vsetvlmax_e32m4();
vfloat32m4_t acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
vfloat32m4_t acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
vfloat32m4_t acc2 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
vfloat32m4_t acc3 = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
size_t m = 0;
while (m < M) {
const size_t vl = __riscv_vsetvl_e32m4(M - m);
// The row term is the same for all four codes, so compute it once.
vuint32m4_t row_off = pq_rvv_kernel::row_offsets(m, vl);
vfloat32m4_t v0 = __riscv_vluxei32_v_f32m4(
sim_table,
pq_rvv_kernel::byte_offsets(row_off, code0, m, vl),
vl);
vfloat32m4_t v1 = __riscv_vluxei32_v_f32m4(
sim_table,
pq_rvv_kernel::byte_offsets(row_off, code1, m, vl),
vl);
vfloat32m4_t v2 = __riscv_vluxei32_v_f32m4(
sim_table,
pq_rvv_kernel::byte_offsets(row_off, code2, m, vl),
vl);
vfloat32m4_t v3 = __riscv_vluxei32_v_f32m4(
sim_table,
pq_rvv_kernel::byte_offsets(row_off, code3, m, vl),
vl);
acc0 = __riscv_vfadd_vv_f32m4_tu(acc0, acc0, v0, vl);
acc1 = __riscv_vfadd_vv_f32m4_tu(acc1, acc1, v1, vl);
acc2 = __riscv_vfadd_vv_f32m4_tu(acc2, acc2, v2, vl);
acc3 = __riscv_vfadd_vv_f32m4_tu(acc3, acc3, v3, vl);
m += vl;
}
result0 = pq_rvv_kernel::reduce(acc0, vlmax);
result1 = pq_rvv_kernel::reduce(acc1, vlmax);
result2 = pq_rvv_kernel::reduce(acc2, vlmax);
result3 = pq_rvv_kernel::reduce(acc3, vlmax);
}
} // namespace pq_code_distance
faiss/impl/scalar_quantizer/sq-rvv.cpp
+5 -69
View File
@@ -12,6 +12,8 @@
#include <faiss/impl/scalar_quantizer/quantizers.h>
#include <faiss/impl/scalar_quantizer/scanners.h>
#include <faiss/impl/scalar_quantizer/similarities.h>
#include <faiss/impl/scalar_quantizer/sq8_coefficients.h>
#include <faiss/impl/scalar_quantizer/sq8_rvv_kernel.h>
#include <riscv_vector.h>
#include <algorithm>
@@ -1967,13 +1969,7 @@ struct DCTemplate<
DCTemplate(size_t d_in, const std::vector<float>& trained)
: d(d_in), a_v(d_in, 0.0f), rmin_v(d_in, 0.0f), e_v(d_in, 0.0f) {
const float* vmin = trained.data();
const float* vdiff = trained.data() + d_in;
for (size_t i = 0; i < d_in; i++) {
float a = vdiff[i] / 255.0f;
a_v[i] = a;
rmin_v[i] = vmin[i] + 0.5f * a;
}
sq8::make_coefficients(trained.data(), d_in, a_v.data(), rmin_v.data());
}
void set_query(const float* x) final {
@@ -1986,61 +1982,7 @@ struct DCTemplate<
/// Full-precision vector L2 over the 1-byte-per-dim code.
float compute_l2(const uint8_t* code) const {
const float* pa = a_v.data();
const float* pe = e_v.data();
// Hoist vsetvl: VLMAX for e8m1 (== f32 lanes per m4 group).
// d==0: guard so vsetvl(0) doesn't return 0 and the prologue
// `if (i + vl <= d)` + chunk loop don't stall.
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
size_t i = 0;
if (i + vl <= d) {
// Software pipeline depth 1: preload the next chunk's code
// bytes at the top of this iteration.
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl); // prologue
i += vl;
for (; i + vl <= d; i += vl) {
vuint8m1_t c8_next = __riscv_vle8_v_u8m1(code + i, vl);
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(c8, vl), vl);
vfloat32m4_t t = __riscv_vle32_v_f32m4(pe + i - vl, vl);
t = __riscv_vfnmsac_vv_f32m4(
t, __riscv_vle32_v_f32m4(pa + i - vl, vl), cf, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
c8 = c8_next; // rotate
}
// Epilogue: process the last full chunk (already loaded).
{
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(c8, vl), vl);
vfloat32m4_t t = __riscv_vle32_v_f32m4(pe + i - vl, vl);
t = __riscv_vfnmsac_vv_f32m4(
t, __riscv_vle32_v_f32m4(pa + i - vl, vl), cf, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
}
}
// Tail: fewer than vl dims left; accumulates into the first
// lanes of acc (safe: reduction below covers all vl lanes).
if (i < d) {
size_t vt = __riscv_vsetvl_e8m1(d - i);
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(c8, vt), vt);
vfloat32m4_t t = __riscv_vle32_v_f32m4(pe + i, vt);
t = __riscv_vfnmsac_vv_f32m4(
t, __riscv_vle32_v_f32m4(pa + i, vt), cf, vt);
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t, t, vt);
}
// Horizontal reduce over all vl lanes
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
return __riscv_vfmv_f_s_f32m1_f32(red);
return sq8_rvv::l2_distance(code, a_v.data(), e_v.data(), d);
}
float query_to_code(const uint8_t* code) const final {
@@ -2112,13 +2054,7 @@ struct DCTemplate<
// vle32 loads stay in bounds (padding lanes never used).
b_v(d_in + 2 * __riscv_vsetvlmax_e8m1(), 0.0f),
k_q(0.0f) {
const float* vmin = trained.data();
const float* vdiff = trained.data() + d_in;
for (size_t i = 0; i < d_in; i++) {
float a = vdiff[i] / 255.0f;
a_v[i] = a;
rmin_v[i] = vmin[i] + 0.5f * a;
}
sq8::make_coefficients(trained.data(), d_in, a_v.data(), rmin_v.data());
}
void set_query(const float* x) final {
faiss/impl/scalar_quantizer/sq8_coefficients.h
+48
View File
@@ -0,0 +1,48 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
/* QT_8bit reconstruction coefficients.
*
* Scalar and portable. A caller on any architecture may include this.
*
* For a QT_8bit non-uniform quantizer with trained values vmin and vdiff:
*
* a[i] = vdiff[i] / 255
* rmin[i] = vmin[i] + 0.5 * a[i]
* e[i] = query[i] - rmin[i]
*
* so reconstruct(code)[i] = rmin[i] + a[i] * code[i].
*/
#include <cstddef>
namespace faiss {
namespace sq8 {
/// Derive the per-dimension scale and offset from a trained quantizer.
/// @param trained 2*d floats: vmin[0..d) then vdiff[0..d)
/// @param levels the largest code value: 255 for QT_8bit, 15 for QT_4bit.
/// Both quantizers reconstruct as rmin[i] + a[i] * code[i],
/// so only the divisor differs.
inline void make_coefficients(
const float* trained,
size_t d,
float* a,
float* rmin,
float levels = 255.0f) {
const float* vmin = trained;
const float* vdiff = trained + d;
for (size_t i = 0; i < d; i++) {
a[i] = vdiff[i] / levels;
rmin[i] = vmin[i] + 0.5f * a[i];
}
}
} // namespace sq8
} // namespace faiss
faiss/impl/scalar_quantizer/sq8_rvv_kernel.h
+546
View File
@@ -0,0 +1,546 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
/* QT_8bit distance kernel for RISC-V Vector.
*
* Free-standing (no exceptions, heap or threads), so a freestanding RISC-V
* target can compile it. sq-rvv.cpp calls `l2_distance`; the block, strided
* and 4-bit forms serve callers that store codes transposed, such as the
* MTIA kernels in faiss/fb/mtia. No Buck target here builds the RISC-V CPU
* path. See sq8_coefficients.h for the meaning of a and e.
*/
#include <riscv_vector.h>
#include <cstddef>
#include <cstdint>
namespace faiss {
namespace sq8_rvv {
/// Candidates per transposed block. One block fills an f32 LMUL 4 group at
/// VLEN=512, so every lane holds one candidate. On a smaller VLEN the block
/// functions loop over the block in pieces of one vector length.
constexpr size_t kBlockVectors = 64;
/// Squared L2 for a whole block, from codes stored transposed: the byte for
/// dimension i of lane c sits at `block[i * kBlockVectors + c]`. Each lane
/// accumulates one candidate, so no reduction is needed.
inline void l2_distance_block(
const uint8_t* block,
const float* a,
const float* e,
size_t d,
float* out) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
for (size_t i = 0; i < d; i++) {
vuint8m1_t c8 =
__riscv_vle8_v_u8m1(block + c + i * kBlockVectors, vl);
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(c8, vl), vl);
// a[i] and e[i] are the same for every candidate in the block.
vfloat32m4_t t = __riscv_vfmv_v_f_f32m4(e[i], vl);
t = __riscv_vfnmsac_vf_f32m4(t, a[i], cf, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
}
__riscv_vse32_v_f32m4(out + c, acc, vl);
c += vl;
}
}
/// Squared L2 for two transposed blocks at once.
///
/// Issuing the loads of both blocks first keeps two loads in flight. Four
/// blocks would need 16 of the 32 vector registers for accumulators and
/// spill.
inline void l2_distance_block_x2(
const uint8_t* block0,
const uint8_t* block1,
const float* a,
const float* e,
size_t d,
float* out0,
float* out1) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vfloat32m4_t acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
vfloat32m4_t acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
for (size_t i = 0; i < d; i++) {
vuint8m1_t v0 =
__riscv_vle8_v_u8m1(block0 + c + i * kBlockVectors, vl);
vuint8m1_t v1 =
__riscv_vle8_v_u8m1(block1 + c + i * kBlockVectors, vl);
const float ei = e[i];
const float ai = a[i];
vfloat32m4_t c0 = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(v0, vl), vl);
vfloat32m4_t t0 = __riscv_vfmv_v_f_f32m4(ei, vl);
t0 = __riscv_vfnmsac_vf_f32m4(t0, ai, c0, vl);
acc0 = __riscv_vfmacc_vv_f32m4(acc0, t0, t0, vl);
vfloat32m4_t c1 = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(v1, vl), vl);
vfloat32m4_t t1 = __riscv_vfmv_v_f_f32m4(ei, vl);
t1 = __riscv_vfnmsac_vf_f32m4(t1, ai, c1, vl);
acc1 = __riscv_vfmacc_vv_f32m4(acc1, t1, t1, vl);
}
__riscv_vse32_v_f32m4(out0 + c, acc0, vl);
__riscv_vse32_v_f32m4(out1 + c, acc1, vl);
c += vl;
}
}
/// The code-dependent inner product term for two transposed blocks at once.
inline void ip_partial_block_x2(
const uint8_t* block0,
const uint8_t* block1,
const float* eq,
size_t d,
float* out0,
float* out1) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vfloat32m4_t acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
vfloat32m4_t acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
for (size_t i = 0; i < d; i++) {
vuint8m1_t v0 =
__riscv_vle8_v_u8m1(block0 + c + i * kBlockVectors, vl);
vuint8m1_t v1 =
__riscv_vle8_v_u8m1(block1 + c + i * kBlockVectors, vl);
const float q = eq[i];
acc0 = __riscv_vfmacc_vf_f32m4(
acc0,
q,
__riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(v0, vl), vl),
vl);
acc1 = __riscv_vfmacc_vf_f32m4(
acc1,
q,
__riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(v1, vl), vl),
vl);
}
__riscv_vse32_v_f32m4(out0 + c, acc0, vl);
__riscv_vse32_v_f32m4(out1 + c, acc1, vl);
c += vl;
}
}
/// The low nibbles of a row of packed QT_4bit codes, as floats.
inline vfloat32m4_t low_nibbles(vuint8m1_t packed, size_t vl) {
return __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(__riscv_vand_vx_u8m1(packed, 0x0F, vl), vl),
vl);
}
/// Squared L2 for a transposed block of QT_4bit codes. Faiss packs
/// dimension i into `code[i / 2]`, even dimensions in the low nibble and odd
/// in the high one, so d dimensions occupy (d + 1) / 2 bytes. For an odd d
/// the last byte holds one dimension, in its low nibble.
inline void l2_distance_4bit_block(
const uint8_t* block,
const float* a,
const float* e,
size_t d,
float* out) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
for (size_t i = 0; i < d / 2; i++) {
vuint8m1_t packed =
__riscv_vle8_v_u8m1(block + c + i * kBlockVectors, vl);
vfloat32m4_t lo = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(
__riscv_vand_vx_u8m1(packed, 0x0F, vl), vl),
vl);
vfloat32m4_t hi = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(
__riscv_vsrl_vx_u8m1(packed, 4, vl), vl),
vl);
vfloat32m4_t t = __riscv_vfmv_v_f_f32m4(e[2 * i], vl);
t = __riscv_vfnmsac_vf_f32m4(t, a[2 * i], lo, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
t = __riscv_vfmv_v_f_f32m4(e[2 * i + 1], vl);
t = __riscv_vfnmsac_vf_f32m4(t, a[2 * i + 1], hi, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
}
if (d % 2 != 0) {
vfloat32m4_t lo = low_nibbles(
__riscv_vle8_v_u8m1(
block + c + (d / 2) * kBlockVectors, vl),
vl);
vfloat32m4_t t = __riscv_vfmv_v_f_f32m4(e[d - 1], vl);
t = __riscv_vfnmsac_vf_f32m4(t, a[d - 1], lo, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
}
__riscv_vse32_v_f32m4(out + c, acc, vl);
c += vl;
}
}
/// The code-dependent inner product term for a transposed QT_4bit block.
inline void ip_partial_4bit_block(
const uint8_t* block,
const float* eq,
size_t d,
float* out) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
for (size_t i = 0; i < d / 2; i++) {
vuint8m1_t packed =
__riscv_vle8_v_u8m1(block + c + i * kBlockVectors, vl);
vfloat32m4_t lo = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(
__riscv_vand_vx_u8m1(packed, 0x0F, vl), vl),
vl);
vfloat32m4_t hi = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(
__riscv_vsrl_vx_u8m1(packed, 4, vl), vl),
vl);
acc = __riscv_vfmacc_vf_f32m4(acc, eq[2 * i], lo, vl);
acc = __riscv_vfmacc_vf_f32m4(acc, eq[2 * i + 1], hi, vl);
}
if (d % 2 != 0) {
vfloat32m4_t lo = low_nibbles(
__riscv_vle8_v_u8m1(
block + c + (d / 2) * kBlockVectors, vl),
vl);
acc = __riscv_vfmacc_vf_f32m4(acc, eq[d - 1], lo, vl);
}
__riscv_vse32_v_f32m4(out + c, acc, vl);
c += vl;
}
}
/// Squared L2 for one QT_4bit candidate in the transposed layout, for the
/// runs an `IDSelector` produces, which start anywhere.
inline float l2_distance_4bit_strided(
const uint8_t* p,
const float* a,
const float* e,
size_t d) {
float acc = 0.0f;
for (size_t i = 0; i < d / 2; i++) {
const uint8_t packed = p[i * kBlockVectors];
const float t0 =
e[2 * i] - a[2 * i] * static_cast<float>(packed & 0x0F);
const float t1 =
e[2 * i + 1] - a[2 * i + 1] * static_cast<float>(packed >> 4);
acc += t0 * t0 + t1 * t1;
}
if (d % 2 != 0) {
const uint8_t packed = p[(d / 2) * kBlockVectors];
const float t = e[d - 1] - a[d - 1] * static_cast<float>(packed & 0x0F);
acc += t * t;
}
return acc;
}
/// The code-dependent inner product term for one strided QT_4bit candidate.
inline float ip_partial_4bit_strided(
const uint8_t* p,
const float* eq,
size_t d) {
float acc = 0.0f;
for (size_t i = 0; i < d / 2; i++) {
const uint8_t packed = p[i * kBlockVectors];
acc += eq[2 * i] * static_cast<float>(packed & 0x0F);
acc += eq[2 * i + 1] * static_cast<float>(packed >> 4);
}
if (d % 2 != 0) {
const uint8_t packed = p[(d / 2) * kBlockVectors];
acc += eq[d - 1] * static_cast<float>(packed & 0x0F);
}
return acc;
}
/// Squared L2 for two transposed QT_4bit blocks at once: half the bytes of
/// QT_8bit and two loads in flight.
inline void l2_distance_4bit_block_x2(
const uint8_t* block0,
const uint8_t* block1,
const float* a,
const float* e,
size_t d,
float* out0,
float* out1) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vfloat32m4_t acc0 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
vfloat32m4_t acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
for (size_t i = 0; i < d / 2; i++) {
vuint8m1_t p0 =
__riscv_vle8_v_u8m1(block0 + c + i * kBlockVectors, vl);
vuint8m1_t p1 =
__riscv_vle8_v_u8m1(block1 + c + i * kBlockVectors, vl);
const float e0 = e[2 * i];
const float e1 = e[2 * i + 1];
const float a0 = a[2 * i];
const float a1 = a[2 * i + 1];
auto step = [&](vuint8m1_t p, vfloat32m4_t acc) {
vfloat32m4_t lo = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(
__riscv_vand_vx_u8m1(p, 0x0F, vl), vl),
vl);
vfloat32m4_t hi = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(
__riscv_vsrl_vx_u8m1(p, 4, vl), vl),
vl);
vfloat32m4_t t = __riscv_vfmv_v_f_f32m4(e0, vl);
t = __riscv_vfnmsac_vf_f32m4(t, a0, lo, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
t = __riscv_vfmv_v_f_f32m4(e1, vl);
t = __riscv_vfnmsac_vf_f32m4(t, a1, hi, vl);
return __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
};
acc0 = step(p0, acc0);
acc1 = step(p1, acc1);
}
if (d % 2 != 0) {
const size_t last = (d / 2) * kBlockVectors;
auto tail = [&](vuint8m1_t p, vfloat32m4_t acc) {
vfloat32m4_t t = __riscv_vfmv_v_f_f32m4(e[d - 1], vl);
t = __riscv_vfnmsac_vf_f32m4(
t, a[d - 1], low_nibbles(p, vl), vl);
return __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
};
acc0 = tail(__riscv_vle8_v_u8m1(block0 + c + last, vl), acc0);
acc1 = tail(__riscv_vle8_v_u8m1(block1 + c + last, vl), acc1);
}
__riscv_vse32_v_f32m4(out0 + c, acc0, vl);
__riscv_vse32_v_f32m4(out1 + c, acc1, vl);
c += vl;
}
}
/// Squared L2 for one candidate whose code is stored transposed: its d bytes
/// sit `kBlockVectors` apart. For runs that do not start on a block
/// boundary, such as the runs an `IDSelector` produces.
inline float l2_distance_strided(
const uint8_t* p,
const float* a,
const float* e,
size_t d) {
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
size_t i = 0;
while (i < d) {
const size_t v = __riscv_vsetvl_e8m1(d - i);
vuint8m1_t c8 = __riscv_vlse8_v_u8m1(
p + i * kBlockVectors,
static_cast<ptrdiff_t>(kBlockVectors),
v);
vfloat32m4_t cf =
__riscv_vfcvt_f_xu_v_f32m4(__riscv_vzext_vf4_u32m4(c8, v), v);
vfloat32m4_t t = __riscv_vle32_v_f32m4(e + i, v);
t = __riscv_vfnmsac_vv_f32m4(t, __riscv_vle32_v_f32m4(a + i, v), cf, v);
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t, t, v);
i += v;
}
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
return __riscv_vfmv_f_s_f32m1_f32(red);
}
/// The code-dependent inner product term for one transposed candidate.
inline float ip_partial_strided(const uint8_t* p, const float* eq, size_t d) {
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
size_t i = 0;
while (i < d) {
const size_t v = __riscv_vsetvl_e8m1(d - i);
vuint8m1_t c8 = __riscv_vlse8_v_u8m1(
p + i * kBlockVectors,
static_cast<ptrdiff_t>(kBlockVectors),
v);
vfloat32m4_t cf =
__riscv_vfcvt_f_xu_v_f32m4(__riscv_vzext_vf4_u32m4(c8, v), v);
acc = __riscv_vfmacc_vv_f32m4_tu(
acc, __riscv_vle32_v_f32m4(eq + i, v), cf, v);
i += v;
}
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
return __riscv_vfmv_f_s_f32m1_f32(red);
}
/// Squared L2 on the 8-bit grid for a transposed block, with integer
/// weights. The host scales the result back.
inline void l2_distance_int_block(
const uint8_t* block,
const uint8_t* qcode,
const int32_t* w,
size_t d,
float* out) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vint32m4_t acc = __riscv_vmv_v_x_i32m4(0, vl);
for (size_t i = 0; i < d; i++) {
vuint8m1_t c8 =
__riscv_vle8_v_u8m1(block + c + i * kBlockVectors, vl);
vint32m4_t c = __riscv_vreinterpret_v_u32m4_i32m4(
__riscv_vzext_vf4_u32m4(c8, vl));
vint32m4_t diff = __riscv_vsub_vx_i32m4(
c, static_cast<int32_t>(qcode[i]), vl);
vint32m4_t sq = __riscv_vmul_vv_i32m4(diff, diff, vl);
acc = __riscv_vmacc_vx_i32m4(acc, w[i], sq, vl);
}
__riscv_vse32_v_f32m4(out + c, __riscv_vfcvt_f_x_v_f32m4(acc, vl), vl);
c += vl;
}
}
/// The code-dependent term of the inner product, for a transposed block.
inline void ip_partial_block(
const uint8_t* block,
const float* eq,
size_t d,
float* out) {
for (size_t c = 0; c < kBlockVectors;) {
const size_t vl = __riscv_vsetvl_e32m4(kBlockVectors - c);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
for (size_t i = 0; i < d; i++) {
vuint8m1_t c8 =
__riscv_vle8_v_u8m1(block + c + i * kBlockVectors, vl);
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(c8, vl), vl);
acc = __riscv_vfmacc_vf_f32m4(acc, eq[i], cf, vl);
}
__riscv_vse32_v_f32m4(out + c, acc, vl);
c += vl;
}
}
/// Squared L2 distance between a query and one 8-bit code.
inline float l2_distance(
const uint8_t* code,
const float* a,
const float* e,
size_t d) {
// VLMAX for e8m1 equals the f32 lane count of an m4 group. Guard d == 0,
// because vsetvl then returns 0 and the chunk loop stalls.
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
size_t i = 0;
if (i + vl <= d) {
// Software pipeline: load the next chunk's code bytes at the top.
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl);
i += vl;
for (; i + vl <= d; i += vl) {
vuint8m1_t c8_next = __riscv_vle8_v_u8m1(code + i, vl);
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
__riscv_vzext_vf4_u32m4(c8, vl), vl);
vfloat32m4_t t = __riscv_vle32_v_f32m4(e + i - vl, vl);
t = __riscv_vfnmsac_vv_f32m4(
t, __riscv_vle32_v_f32m4(a + i - vl, vl), cf, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
c8 = c8_next;
}
vfloat32m4_t cf =
__riscv_vfcvt_f_xu_v_f32m4(__riscv_vzext_vf4_u32m4(c8, vl), vl);
vfloat32m4_t t = __riscv_vle32_v_f32m4(e + i - vl, vl);
t = __riscv_vfnmsac_vv_f32m4(
t, __riscv_vle32_v_f32m4(a + i - vl, vl), cf, vl);
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
}
// Tail: fewer than vl dimensions remain. The reduction covers every lane.
if (i < d) {
const size_t vt = __riscv_vsetvl_e8m1(d - i);
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
vfloat32m4_t cf =
__riscv_vfcvt_f_xu_v_f32m4(__riscv_vzext_vf4_u32m4(c8, vt), vt);
vfloat32m4_t t = __riscv_vle32_v_f32m4(e + i, vt);
t = __riscv_vfnmsac_vv_f32m4(
t, __riscv_vle32_v_f32m4(a + i, vt), cf, vt);
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t, t, vt);
}
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
return __riscv_vfmv_f_s_f32m1_f32(red);
}
/// Inner product between a query and one 8-bit code, without the constant.
///
/// Expanding the reconstruction,
///
/// sum q[i] * v[i] = sum q[i]*rmin[i] + sum (q[i]*a[i]) * code[i]
///
/// the first term does not depend on the code, so the caller adds it. This
/// computes the second, with u[i] = q[i] * a[i].
inline float ip_partial(const uint8_t* code, const float* u, size_t d) {
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
size_t i = 0;
for (; i + vl <= d; i += vl) {
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl);
vfloat32m4_t cf =
__riscv_vfcvt_f_xu_v_f32m4(__riscv_vzext_vf4_u32m4(c8, vl), vl);
acc = __riscv_vfmacc_vv_f32m4(
acc, cf, __riscv_vle32_v_f32m4(u + i, vl), vl);
}
if (i < d) {
const size_t vt = __riscv_vsetvl_e8m1(d - i);
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
vfloat32m4_t cf =
__riscv_vfcvt_f_xu_v_f32m4(__riscv_vzext_vf4_u32m4(c8, vt), vt);
acc = __riscv_vfmacc_vv_f32m4_tu(
acc, cf, __riscv_vle32_v_f32m4(u + i, vt), vt);
}
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
return __riscv_vfmv_f_s_f32m1_f32(red);
}
/// Squared L2 distance with the query quantized onto the same 8-bit grid.
///
/// sum (q[i] - v[i])^2 = sum a[i]^2 * (qcode[i] - code[i])^2
///
/// The weights hold a[i]^2 as integers, so the loop is integer only. The
/// caller scales the result back and must keep max(w) * 65025 * d inside
/// int32.
inline float l2_distance_int(
const uint8_t* code,
const uint8_t* qcode,
const int32_t* w,
size_t d) {
// Accumulate at LMUL 4, not 8. A group of eight uses eight of the 32
// vector registers and leaves only four groups, which is slower.
vint32m4_t sum = __riscv_vmv_v_x_i32m4(0, __riscv_vsetvlmax_e32m4());
size_t i = 0;
while (i < d) {
const size_t vl = __riscv_vsetvl_e16m2(d - i);
vint16m2_t c = __riscv_vreinterpret_v_u16m2_i16m2(
__riscv_vzext_vf2_u16m2(__riscv_vle8_v_u8m1(code + i, vl), vl));
vint16m2_t q =
__riscv_vreinterpret_v_u16m2_i16m2(__riscv_vzext_vf2_u16m2(
__riscv_vle8_v_u8m1(qcode + i, vl), vl));
vint16m2_t diff = __riscv_vsub_vv_i16m2(q, c, vl);
// A difference of two bytes fits in int16 and its square in int32.
vint32m4_t sq = __riscv_vwmul_vv_i32m4(diff, diff, vl);
sum = __riscv_vmacc_vv_i32m4_tu(
sum, sq, __riscv_vle32_v_i32m4(w + i, vl), vl);
i += vl;
}
vint32m1_t r = __riscv_vredsum_vs_i32m4_i32m1(
sum, __riscv_vmv_v_x_i32m1(0, 1), __riscv_vsetvlmax_e32m4());
return static_cast<float>(__riscv_vmv_x_s_i32m1_i32(r));
}
} // namespace sq8_rvv
} // namespace faiss
faiss/utils/simd_impl/distances_rvv.cpp
+70 -35
View File
@@ -14,10 +14,19 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/extra_distances.h>
#include <riscv_vector.h>
#include <faiss/utils/simd_impl/distances_rvv_kernel.h>
#include <vector>
namespace faiss {
// Clang rejects a builtin used as a function argument, so wrap the call in a
// lambda. GCC accepts the bare name; the wrapper compiles on both.
#define FAISS_RVV_REDUCE_OP(intrinsic) \
[](auto value_, auto init_, size_t vl_) { \
return intrinsic(value_, init_, vl_); \
}
template <typename Vec, typename Reduce>
static inline float rvv_reduce(
Vec value,
@@ -40,7 +49,10 @@ static inline size_t rvv_argmin(const float* values, size_t n) {
i += vl;
}
float min_val = rvv_reduce(
vmin, vlmax, __builtin_inff(), __riscv_vfredmin_vs_f32m8_f32m1);
vmin,
vlmax,
__builtin_inff(),
FAISS_RVV_REDUCE_OP(__riscv_vfredmin_vs_f32m8_f32m1));
i = 0;
while (i < n) {
size_t vl = __riscv_vsetvl_e32m8(n - i);
@@ -185,7 +197,11 @@ float fvec_norm_L2sqr<SIMDLevel::RISCV_RVV>(const float* x, size_t d) {
acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
i += vl;
}
return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
return rvv_reduce(
acc,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m8_f32m1));
}
template <>
@@ -193,18 +209,7 @@ float fvec_L2sqr<SIMDLevel::RISCV_RVV>(
const float* x,
const float* y,
size_t 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;
}
return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
return rvv_kernel::l2_sqr(x, y, d);
}
template <>
@@ -212,17 +217,7 @@ float fvec_inner_product<SIMDLevel::RISCV_RVV>(
const float* x,
const float* y,
size_t 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;
}
return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
return rvv_kernel::inner_product(x, y, d);
}
template <>
@@ -239,7 +234,11 @@ float fvec_L1<SIMDLevel::RISCV_RVV>(const float* x, const float* y, size_t d) {
acc = __riscv_vfadd_vv_f32m8_tu(acc, acc, vx, vl);
i += vl;
}
return rvv_reduce(acc, vlmax, 0.0f, __riscv_vfredusum_vs_f32m8_f32m1);
return rvv_reduce(
acc,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m8_f32m1));
}
template <>
@@ -259,7 +258,11 @@ float fvec_Linf<SIMDLevel::RISCV_RVV>(
vmax = __riscv_vfmax_vv_f32m8_tu(vmax, vmax, vx, vl);
i += vl;
}
return rvv_reduce(vmax, vlmax, 0.0f, __riscv_vfredmax_vs_f32m8_f32m1);
return rvv_reduce(
vmax,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredmax_vs_f32m8_f32m1));
}
template <>
@@ -293,10 +296,26 @@ void fvec_inner_product_batch_4<SIMDLevel::RISCV_RVV>(
vacc3 = __riscv_vfmacc_vv_f32m4_tu(vacc3, vx, vy, vl);
i += vl;
}
dis0 = rvv_reduce(vacc0, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis1 = rvv_reduce(vacc1, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis2 = rvv_reduce(vacc2, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis3 = rvv_reduce(vacc3, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis0 = rvv_reduce(
vacc0,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
dis1 = rvv_reduce(
vacc1,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
dis2 = rvv_reduce(
vacc2,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
dis3 = rvv_reduce(
vacc3,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
}
template <>
@@ -334,10 +353,26 @@ void fvec_L2sqr_batch_4<SIMDLevel::RISCV_RVV>(
vacc3 = __riscv_vfmacc_vv_f32m4_tu(vacc3, vy, vy, vl);
i += vl;
}
dis0 = rvv_reduce(vacc0, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis1 = rvv_reduce(vacc1, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis2 = rvv_reduce(vacc2, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis3 = rvv_reduce(vacc3, vlmax, 0.0f, __riscv_vfredusum_vs_f32m4_f32m1);
dis0 = rvv_reduce(
vacc0,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
dis1 = rvv_reduce(
vacc1,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
dis2 = rvv_reduce(
vacc2,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
dis3 = rvv_reduce(
vacc3,
vlmax,
0.0f,
FAISS_RVV_REDUCE_OP(__riscv_vfredusum_vs_f32m4_f32m1));
}
template <>
faiss/utils/simd_impl/distances_rvv_kernel.h
+68
View File
@@ -0,0 +1,68 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
/* Vector-to-vector distance kernels for RISC-V Vector.
*
* Free-standing (no exceptions, heap or threads), so a freestanding RISC-V
* target can compile it. distances_rvv.cpp calls these; no Buck target here
* builds that RISC-V CPU path.
*
* These inline the horizontal reduction rather than passing the intrinsic to
* a helper. Clang rejects a builtin used as a function argument.
*
* They accumulate at LMUL 4, not 8: a group of eight leaves only four
* register groups, which is slower in a scan loop. The order of the
* additions differs from LMUL 8, so the last bits can differ.
*/
#include <cstddef>
#include <riscv_vector.h>
namespace faiss {
namespace rvv_kernel {
/// Squared L2 distance between two fp32 vectors.
inline float l2_sqr(const float* x, const float* y, size_t d) {
const size_t vlmax = __riscv_vsetvlmax_e32m4();
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
size_t i = 0;
while (i < d) {
const size_t vl = __riscv_vsetvl_e32m4(d - i);
vfloat32m4_t vx = __riscv_vle32_v_f32m4(x + i, vl);
vfloat32m4_t vy = __riscv_vle32_v_f32m4(y + i, vl);
vx = __riscv_vfsub_vv_f32m4(vx, vy, vl);
acc = __riscv_vfmacc_vv_f32m4_tu(acc, vx, vx, vl);
i += vl;
}
vfloat32m1_t init = __riscv_vfmv_s_f_f32m1(0.0f, 1);
vfloat32m1_t r = __riscv_vfredusum_vs_f32m4_f32m1(acc, init, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(r);
}
/// Inner product of two fp32 vectors.
inline float inner_product(const float* x, const float* y, size_t d) {
const size_t vlmax = __riscv_vsetvlmax_e32m4();
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vlmax);
size_t i = 0;
while (i < d) {
const size_t vl = __riscv_vsetvl_e32m4(d - i);
vfloat32m4_t vx = __riscv_vle32_v_f32m4(x + i, vl);
vfloat32m4_t vy = __riscv_vle32_v_f32m4(y + i, vl);
acc = __riscv_vfmacc_vv_f32m4_tu(acc, vx, vy, vl);
i += vl;
}
vfloat32m1_t init = __riscv_vfmv_s_f_f32m1(0.0f, 1);
vfloat32m1_t r = __riscv_vfredusum_vs_f32m4_f32m1(acc, init, vlmax);
return __riscv_vfmv_f_s_f32m1_f32(r);
}
} // namespace rvv_kernel
} // namespace faiss