Summary:
- FastScan's pq4 kernel uses vector abstractions that are 256-bit wide - e.g. `simd32uint8` and `simd16uint16`
- These are performant for processors whose SIMD registers are 256-bits
- However, on Arm, these are implemented using 2 128-bit NEON registers - which leads to register pressure
- For example, the original algorithm which accumulates NQ queries x 32 db elements uses these accumulators: `simd16uint16 accu[NQA][4]` which add up to 24 Neon registers (for the common case where NQA=3) - out of a total of 32 avaliable NEON registers.
- With other temp registers for code vectors and LUT, the compiler spills registers to stack
This change:
- Adds a native NEON kernel equivalent for `pq4_kernel_qbs_256` which accumulates NQ queries for a single basic block of 32 database vectors (it uses 12 accumulators for NQA=3) --- (we'll raise another PR for speeding up `kernel_accumulate_block` which accumulates NQ queries for multiple 32-vectoer DB blocks.)
- keeps the existing QBS decomposition, PQ code packing, LUT packing and the PQ4 accumulate loop - this is now shared between the neon kernel and the existing simd256 kernel
- removes 126 Lines of deadcode from `decompose_qbs.cpp` - please let me know if these were left for a reason.
## Tests:
This change is covered by existing tests: `test_fast_scan`, `test_fast_scan_ivf`, `test_sq_fastscan`, `test_ivf_sq_fastscan` - all these tests path.
Also measured R@1, R@10 and R@100 for index types below for baseline and this PR and made sure they're identical.
## End to End Performance:
with topk=10, nprobe=64, Neoverse-V3
R@1, R@10 and R@100 match exactly with baseline
| Dataset | Index | Threads | Speedup |
|---|---|---:|---:|
| sift1M | `PQ32x4fs` | 32 | **1.412x** |
| sift1M | `PQ32x4fs` | 96 | **1.415x** |
| sift1M | `PQ64x4fs` | 32 | **1.353x** |
| sift1M | `PQ64x4fs` | 96 | **1.351x** |
| sift1M | `PQ128x4fs` | 32 | **1.310x** |
| sift1M | `PQ128x4fs` | 96 | **1.316x** |
| sift1M | `IVF4096,PQ32x4fs` | 32 | **1.104x** |
| sift1M | `IVF4096,PQ32x4fs` | 96 | **1.114x** |
| sift1M | `IVF4096,PQ64x4fs` | 32 | **1.132x** |
| sift1M | `IVF4096,PQ64x4fs` | 96 | **1.144x** |
| sift1M | `IVF4096,PQ128x4fs` | 32 | **1.153x** |
| sift1M | `IVF4096,PQ128x4fs` | 96 | 1.095x |
| sift1M | `IVF4096,SQ4fs` | 32 | **1.156x** |
| sift1M | `IVF4096,SQ4fs` | 96 | 1.082x |
| sift1M | `IVF4096,RQ30x4fs_Nrq2x4` | 32 | **1.102x** |
| sift1M | `IVF4096,RQ30x4fs_Nrq2x4` | 96 | 1.047x |
| sift1M | `IVF4096,RaBitQfs` | 32 | 1.023x |
| sift1M | `IVF4096,RaBitQfs` | 96 | 1.007x |
| bigann10M | `IVF16384,PQ128x4fs` | 32 | 1.087x |
| bigann10M | `IVF16384,PQ128x4fs` | 96 | 1.007x |
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5709
Reviewed By: junjieqi
Differential Revision: D123883104
Pulled By: mnorris11
fbshipit-source-id: a97fa22b42c5ea238cc4220190a4e775a0a30eac
Fadi Arafeh
a2f7943005