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
Faiss
Faiss is a library for efficient similarity search and clustering of dense vectors. It contains algorithms that search in sets of vectors of any size, up to ones that possibly do not fit in RAM. It also contains supporting code for evaluation and parameter tuning. Faiss is written in C++ with complete wrappers for Python/numpy. Some of the most useful algorithms are implemented on the GPU. It is developed primarily at Meta's Fundamental AI Research group.
News
See CHANGELOG.md for detailed information about latest features.
Introduction
Faiss contains several methods for similarity search. It assumes that the instances are represented as vectors and are identified by an integer, and that the vectors can be compared with L2 (Euclidean) distances or dot products. Vectors that are similar to a query vector are those that have the lowest L2 distance or the highest dot product with the query vector. It also supports cosine similarity, since this is a dot product on normalized vectors.
Some of the methods, like those based on binary vectors and compact quantization codes, solely use a compressed representation of the vectors and do not require to keep the original vectors. This generally comes at the cost of a less precise search but these methods can scale to billions of vectors in main memory on a single server. Other methods, like HNSW and NSG add an indexing structure on top of the raw vectors to make searching more efficient.
The GPU implementation can accept input from either CPU or GPU memory. On a server with GPUs, the GPU indexes can be used a drop-in replacement for the CPU indexes (e.g., replace IndexFlatL2 with GpuIndexFlatL2) and copies to/from GPU memory are handled automatically. Results will be faster however if both input and output remain resident on the GPU. Both single and multi-GPU usage is supported.
Installing
Faiss comes with precompiled libraries for Anaconda in Python, see faiss-cpu, faiss-gpu and faiss-gpu-cuvs. The library is mostly implemented in C++, the only dependency is a BLAS implementation. Optional GPU support is provided via CUDA or AMD ROCm, and the Python interface is also optional. The backend GPU implementations of NVIDIA cuVS can also be enabled optionally. It compiles with cmake. See INSTALL.md for details.
How Faiss works
Faiss is built around an index type that stores a set of vectors, and provides a function to search in them with L2 and/or dot product vector comparison. Some index types are simple baselines, such as exact search. Most of the available indexing structures correspond to various trade-offs with respect to
- search time
- search quality
- memory used per index vector
- training time
- adding time
- need for external data for unsupervised training
The optional GPU implementation provides what is likely (as of March 2017) the fastest exact and approximate (compressed-domain) nearest neighbor search implementation for high-dimensional vectors, fastest Lloyd's k-means, and fastest small k-selection algorithm known. The implementation is detailed here.
Full documentation of Faiss
The following are entry points for documentation:
- the full documentation can be found on the wiki page, including a tutorial, a FAQ and a troubleshooting section
- the doxygen documentation gives per-class information extracted from code comments
- to reproduce results from our research papers, Polysemous codes and Billion-scale similarity search with GPUs, refer to the benchmarks README. For Link and code: Fast indexing with graphs and compact regression codes, see the link_and_code README
Authors
The main authors of Faiss are:
- Hervé Jégou initiated the Faiss project and wrote its first implementation
- Matthijs Douze implemented most of the CPU Faiss
- Jeff Johnson implemented all of the GPU Faiss
- Lucas Hosseini implemented the binary indexes and the build system
- Chengqi Deng implemented NSG, NNdescent and much of the additive quantization code.
- Alexandr Guzhva many optimizations: SIMD, memory allocation and layout, fast decoding kernels for vector codecs, etc.
- Gergely Szilvasy build system, benchmarking framework.
Reference
References to cite when you use Faiss in a research paper:
@article{douze2024faiss,
title={The Faiss library},
author={Matthijs Douze and Alexandr Guzhva and Chengqi Deng and Jeff Johnson and Gergely Szilvasy and Pierre-Emmanuel Mazaré and Maria Lomeli and Lucas Hosseini and Hervé Jégou},
year={2024},
eprint={2401.08281},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
For the GPU version of Faiss, please cite:
@article{johnson2019billion,
title={Billion-scale similarity search with {GPUs}},
author={Johnson, Jeff and Douze, Matthijs and J{\'e}gou, Herv{\'e}},
journal={IEEE Transactions on Big Data},
volume={7},
number={3},
pages={535--547},
year={2019},
publisher={IEEE}
}
Join the Faiss community
For public discussion of Faiss or for questions, visit https://github.com/facebookresearch/faiss/discussions.
We monitor the issues page of the repository. You can report bugs, ask questions, etc.
Legal
Faiss is MIT-licensed, refer to the LICENSE file in the top level directory.
Copyright © Meta Platforms, Inc.