mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Add RVV implementation for decode_bf16_simd (#5641)
Summary: Add a RISC-V Vector (RVV) implementation to `decode_bf16_simd()`. The RVV path preserves the existing bit-level conversion semantics: 1. load BF16 values as unsigned 16-bit integers; 2. zero-extend them to unsigned 32-bit integers; 3. shift left by 16 bits; 4. reinterpret the resulting bits as FP32. Inputs shorter than 16 elements retain the scalar path to avoid short-vector overhead. Larger inputs use fixed-VL RVV chunks followed by a runtime-VL tail. ## Validation Tested on native RISC-V hardware with GCC 15.3.0 and `-march=rv64gcv -mabi=lp64d`, based on Faiss commit `b5a14632a6f2bfdc824c86d683b9d6e9205c2604`. Correctness: - exhaustively tested all 65,536 BF16 bit patterns; - tested 31 input sizes, including zero, short inputs, vector boundaries, tails, and large arrays; - checked guard elements for output overwrites; - zero correctness failures. Representative microbenchmark speedups over the scalar implementation: | Elements | Speedup | |---:|---:| | 7 | 1.01x | | 16 | 1.58x | | 31 | 2.43x | | 64 | 3.46x | | 128 | 3.88x | | 1024 | 4.38x | | 4096 | 4.37x | | 65536 | 4.30x | ## Scope This change only modifies `faiss/utils/bf16.h`. It is independent of the active ScalarQuantizer RVV changes in PRs https://github.com/facebookresearch/faiss/issues/5535 and https://github.com/facebookresearch/faiss/issues/5539. ## Notes - The `n == 1` fast path and the `n < 16` scalar loop exist because at VLEN=128 an `e16m1` chunk covers only 8 elements, and for very short inputs the `vsetvl`/load/extend/store sequence costs more than the equivalent scalar work. The measured crossover is at 16 elements. - The bulk loop uses a hoisted `vsetvlmax_e16m1()` so `vsetvl` is not re-executed per iteration; only the final partial chunk calls `vsetvl` again. - `vuint16m1` → `vuint32m2` uses `vzext.vf2` (LMUL widening), so the 32-bit shift and the FP32 store run at `m2`. This keeps register pressure low while doubling the elements per chunk. - No rounding or saturation is involved: BF16 → FP32 is an exact bit widening, so the vector path is bit-identical to the scalar path for every input, including NaN payloads and subnormals. - The guard is `#elif defined(__riscv_vector)` inside the existing arch dispatch, so the AVX2/AVX512 paths are untouched and a build without RVV falls through to the scalar loop unchanged. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5641 Reviewed By: alibeklfc Differential Revision: D121445956 Pulled By: mnorris11 fbshipit-source-id: 4b4ec0bad6473030849ae9275cbdad0bc3e087ee Co-authored-by: ihb2032 <[email protected]> Co-authored-by: lyd1992 <[email protected]> Co-authored-by: Yuansheng <[email protected]>
This commit is contained in:
committed by
meta-codesync[bot]
co-authored by
ihb2032
lyd1992
Yuansheng
parent
fdb9535c15
commit
6c8378e4f1
@@ -14,6 +14,10 @@
|
|||||||
#include <immintrin.h>
|
#include <immintrin.h>
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
#if defined(__riscv_vector)
|
||||||
|
#include <riscv_vector.h>
|
||||||
|
#endif
|
||||||
|
|
||||||
namespace faiss {
|
namespace faiss {
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
@@ -61,6 +65,34 @@ inline void decode_bf16_simd(const uint16_t* src, float* dst, size_t n) {
|
|||||||
w = _mm512_slli_epi32(w, 16);
|
w = _mm512_slli_epi32(w, 16);
|
||||||
_mm512_storeu_ps(dst + i, _mm512_castsi512_ps(w));
|
_mm512_storeu_ps(dst + i, _mm512_castsi512_ps(w));
|
||||||
}
|
}
|
||||||
|
#elif defined(__riscv_vector)
|
||||||
|
if (n == 1) {
|
||||||
|
dst[0] = decode_bf16(src[0]);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if (n < 16) {
|
||||||
|
for (; i < n; ++i) {
|
||||||
|
dst[i] = decode_bf16(src[i]);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
const size_t vlmax = __riscv_vsetvlmax_e16m1();
|
||||||
|
for (; i + vlmax <= n; i += vlmax) {
|
||||||
|
vuint16m1_t v = __riscv_vle16_v_u16m1(src + i, vlmax);
|
||||||
|
vuint32m2_t w = __riscv_vzext_vf2_u32m2(v, vlmax);
|
||||||
|
w = __riscv_vsll_vx_u32m2(w, 16, vlmax);
|
||||||
|
__riscv_vse32_v_f32m2(
|
||||||
|
dst + i, __riscv_vreinterpret_v_u32m2_f32m2(w), vlmax);
|
||||||
|
}
|
||||||
|
if (i < n) {
|
||||||
|
const size_t vl = __riscv_vsetvl_e16m1(n - i);
|
||||||
|
vuint16m1_t v = __riscv_vle16_v_u16m1(src + i, vl);
|
||||||
|
vuint32m2_t w = __riscv_vzext_vf2_u32m2(v, vl);
|
||||||
|
w = __riscv_vsll_vx_u32m2(w, 16, vl);
|
||||||
|
__riscv_vse32_v_f32m2(
|
||||||
|
dst + i, __riscv_vreinterpret_v_u32m2_f32m2(w), vl);
|
||||||
|
i += vl;
|
||||||
|
}
|
||||||
#endif
|
#endif
|
||||||
for (; i < n; i++) {
|
for (; i < n; i++) {
|
||||||
dst[i] = decode_bf16(src[i]);
|
dst[i] = decode_bf16(src[i]);
|
||||||
|
|||||||
Reference in New Issue
Block a user