diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 058fc0fba..513c5dc11 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -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 diff --git a/faiss/impl/pq_code_distance/pq_rvv_kernel.h b/faiss/impl/pq_code_distance/pq_rvv_kernel.h new file mode 100644 index 000000000..133e576fa --- /dev/null +++ b/faiss/impl/pq_code_distance/pq_rvv_kernel.h @@ -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 +#include + +#include + +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 diff --git a/faiss/impl/pq_code_distance/rvv.cpp b/faiss/impl/pq_code_distance/rvv.cpp index 8bad873f7..94dee5590 100644 --- a/faiss/impl/pq_code_distance/rvv.cpp +++ b/faiss/impl/pq_code_distance/rvv.cpp @@ -7,12 +7,15 @@ #ifdef COMPILE_SIMD_RISCV_RVV +#include + #include +#include 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( size_t M, const float* sim_table, const uint8_t* code) { - return PQCodeDistanceScalar::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( float& result1, float& result2, float& result3) { - PQCodeDistanceScalar::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 diff --git a/faiss/impl/scalar_quantizer/sq-rvv.cpp b/faiss/impl/scalar_quantizer/sq-rvv.cpp index 1098c2b05..e891d77d1 100644 --- a/faiss/impl/scalar_quantizer/sq-rvv.cpp +++ b/faiss/impl/scalar_quantizer/sq-rvv.cpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include #include #include @@ -1967,13 +1969,7 @@ struct DCTemplate< DCTemplate(size_t d_in, const std::vector& 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 { diff --git a/faiss/impl/scalar_quantizer/sq8_coefficients.h b/faiss/impl/scalar_quantizer/sq8_coefficients.h new file mode 100644 index 000000000..500e4f199 --- /dev/null +++ b/faiss/impl/scalar_quantizer/sq8_coefficients.h @@ -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 + +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 diff --git a/faiss/impl/scalar_quantizer/sq8_rvv_kernel.h b/faiss/impl/scalar_quantizer/sq8_rvv_kernel.h new file mode 100644 index 000000000..51143b806 --- /dev/null +++ b/faiss/impl/scalar_quantizer/sq8_rvv_kernel.h @@ -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 + +#include +#include + +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(packed & 0x0F); + const float t1 = + e[2 * i + 1] - a[2 * i + 1] * static_cast(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(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(packed & 0x0F); + acc += eq[2 * i + 1] * static_cast(packed >> 4); + } + if (d % 2 != 0) { + const uint8_t packed = p[(d / 2) * kBlockVectors]; + acc += eq[d - 1] * static_cast(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(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(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(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(__riscv_vmv_x_s_i32m1_i32(r)); +} + +} // namespace sq8_rvv +} // namespace faiss diff --git a/faiss/utils/simd_impl/distances_rvv.cpp b/faiss/utils/simd_impl/distances_rvv.cpp index dc2d872dc..5b18856fb 100644 --- a/faiss/utils/simd_impl/distances_rvv.cpp +++ b/faiss/utils/simd_impl/distances_rvv.cpp @@ -14,10 +14,19 @@ #include #include #include + +#include #include 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 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(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( 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( 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(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( 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( 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( 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 <> diff --git a/faiss/utils/simd_impl/distances_rvv_kernel.h b/faiss/utils/simd_impl/distances_rvv_kernel.h new file mode 100644 index 000000000..4cd429924 --- /dev/null +++ b/faiss/utils/simd_impl/distances_rvv_kernel.h @@ -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 + +#include + +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