From a2f79430051d5c8c458300645a15c6efa496d1f4 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh Date: Wed, 7 Oct 2026 13:10:01 -0700 Subject: [PATCH] Optimize 4-bit PQ fastscan accumulation on Arm CPUs (#5709) 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 --- faiss/CMakeLists.txt | 8 +- faiss/impl/fast_scan/accumulate_loops.h | 123 ++++++++-------- faiss/impl/fast_scan/accumulate_loops_512.h | 14 +- faiss/impl/fast_scan/decompose_qbs.h | 141 ++---------------- faiss/impl/fast_scan/dispatching.h | 24 ++- faiss/impl/fast_scan/kernels_neon.h | 155 ++++++++++++++++++++ faiss/impl/fast_scan/kernels_simd256.h | 23 ++- tests/test_single_query_collect_handler.cpp | 14 +- 8 files changed, 296 insertions(+), 206 deletions(-) create mode 100644 faiss/impl/fast_scan/kernels_neon.h diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 155995098..db2eb0c10 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -1,4 +1,9 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# Portions 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. + +# Portions Copyright 2026 Arm Limited and/or its affiliates # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. @@ -357,6 +362,7 @@ set(FAISS_HEADERS impl/fast_scan/dispatching.h impl/fast_scan/fast_scan.h impl/fast_scan/decompose_qbs.h + impl/fast_scan/kernels_neon.h impl/fast_scan/kernels_simd256.h impl/fast_scan/kernels_simd512.h impl/fast_scan/rabitq_dispatching.h diff --git a/faiss/impl/fast_scan/accumulate_loops.h b/faiss/impl/fast_scan/accumulate_loops.h index 34d5e174a..9c640b0e3 100644 --- a/faiss/impl/fast_scan/accumulate_loops.h +++ b/faiss/impl/fast_scan/accumulate_loops.h @@ -1,5 +1,13 @@ /* - * Copyright (c) Meta Platforms, Inc. and affiliates. + * Portions 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. + */ + +/* + * Portions Copyright 2026 Arm Limited and/or its affiliates + * * * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. @@ -14,22 +22,14 @@ * Contains: * - accumulate_fixed_blocks / pq4_accumulate_loop_fixed_scaler * (search_1 multi-BB path, bbs > 32) - * - accumulate_q_4step_256 / pq4_accumulate_loop_qbs_fixed_scaler_256 - * (QBS path, bbs == 32, 256-bit kernel only) + * - SIMD-level-selected QBS traversal (bbs == 32) * - * The QBS helpers here use pq4_kernel_qbs_256 exclusively (not - * decompose_qbs.h) because decompose_qbs.h includes kernels_simd512.h - * whose 512-bit types need explicit SIMD levels. The 512-bit QBS path - * lives in accumulate_loops_512.h, used by the AVX512 per-ISA TU. + * The default PQ4QBSKernel uses the 256-bit kernel. Architecture-specific + * headers may specialize it while reusing the traversal and QBS dispatch. * * All functions live in `namespace faiss` (not anonymous) so they can be * shared by both the per-SIMD TU dispatcher (dispatching.h) and the old * free-function search paths (pq4_fast_scan_search_1.cpp). - * - * The QBS helpers here always use pq4_kernel_qbs_256 (never 512-bit). - * This is required for the per-SIMD DD TUs where SINGLE_SIMD_LEVEL=NONE - * leaves 512-bit types empty. The old pq4_fast_scan_search_qbs.cpp - * continues to use decompose_qbs.h which includes both 256 and 512 paths. */ #include @@ -37,6 +37,9 @@ #include #include #include +#ifdef COMPILE_SIMD_ARM_NEON +#include +#endif #include namespace faiss { @@ -112,15 +115,11 @@ void pq4_accumulate_loop_fixed_scaler( } /*************************************************************** - * QBS path helpers (bbs == 32, 256-bit kernel only) + * QBS path helpers (bbs == 32) ***************************************************************/ -template < - int QBS, - SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, - class ResultHandler, - class Scaler> -void accumulate_q_4step_256( +template +void accumulate_q_4step_for_simd( size_t ntotal2, int nsq, const uint8_t* codes, @@ -135,33 +134,30 @@ void accumulate_q_4step_256( constexpr int SQ = Q1 + Q2 + Q3 + Q4; for_each_block<32>(ntotal2, codes, block_stride, res, [&](size_t) { - FixedStorageHandler res2; + FixedStorageHandler res2; const uint8_t* LUT = LUT0; - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + PQ4QBSKernel::template run(nsq, codes, LUT, res2, scaler); LUT += Q1 * nsq * 16; if (Q2 > 0) { res2.set_block_origin(Q1, 0); - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + PQ4QBSKernel::template run(nsq, codes, LUT, res2, scaler); LUT += Q2 * nsq * 16; } if (Q3 > 0) { res2.set_block_origin(Q1 + Q2, 0); - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + PQ4QBSKernel::template run(nsq, codes, LUT, res2, scaler); LUT += Q3 * nsq * 16; } if (Q4 > 0) { res2.set_block_origin(Q1 + Q2 + Q3, 0); - pq4_kernel_qbs_256(nsq, codes, LUT, res2, scaler); + PQ4QBSKernel::template run(nsq, codes, LUT, res2, scaler); } res2.to_other_handler(res); }); } -template < - SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, - class ResultHandler, - class Scaler> -void pq4_accumulate_loop_qbs_fixed_scaler_256( +template +void pq4_accumulate_loop_qbs_fixed_scaler_simd( int qbs, size_t ntotal2, int nsq, @@ -175,38 +171,37 @@ void pq4_accumulate_loop_qbs_fixed_scaler_256( assert(is_aligned_pointer(LUT0)); switch (qbs) { -#define FAISS_QBS256_DISPATCH(QBS) \ +#define FAISS_QBS_DISPATCH(QBS) \ case QBS: \ - accumulate_q_4step_256( \ + accumulate_q_4step_for_simd( \ ntotal2, nsq, codes, LUT0, res, scaler, block_stride); \ return; - FAISS_QBS256_DISPATCH(0x3333); // 12 - FAISS_QBS256_DISPATCH(0x2333); // 11 - FAISS_QBS256_DISPATCH(0x2233); // 10 - FAISS_QBS256_DISPATCH(0x333); // 9 - FAISS_QBS256_DISPATCH(0x2223); // 9 - FAISS_QBS256_DISPATCH(0x233); // 8 - FAISS_QBS256_DISPATCH(0x1223); // 8 - FAISS_QBS256_DISPATCH(0x223); // 7 - FAISS_QBS256_DISPATCH(0x34); // 7 - FAISS_QBS256_DISPATCH(0x133); // 7 - FAISS_QBS256_DISPATCH(0x6); // 6 - FAISS_QBS256_DISPATCH(0x33); // 6 - FAISS_QBS256_DISPATCH(0x123); // 6 - FAISS_QBS256_DISPATCH(0x222); // 6 - FAISS_QBS256_DISPATCH(0x23); // 5 - FAISS_QBS256_DISPATCH(0x5); // 5 - FAISS_QBS256_DISPATCH(0x13); // 4 - FAISS_QBS256_DISPATCH(0x22); // 4 - FAISS_QBS256_DISPATCH(0x4); // 4 - FAISS_QBS256_DISPATCH(0x3); // 3 - FAISS_QBS256_DISPATCH(0x21); // 3 - FAISS_QBS256_DISPATCH(0x2); // 2 - FAISS_QBS256_DISPATCH(0x1); // 1 -#undef FAISS_QBS256_DISPATCH + FAISS_QBS_DISPATCH(0x3333); // 12 + FAISS_QBS_DISPATCH(0x2333); // 11 + FAISS_QBS_DISPATCH(0x2233); // 10 + FAISS_QBS_DISPATCH(0x333); // 9 + FAISS_QBS_DISPATCH(0x2223); // 9 + FAISS_QBS_DISPATCH(0x233); // 8 + FAISS_QBS_DISPATCH(0x1223); // 8 + FAISS_QBS_DISPATCH(0x223); // 7 + FAISS_QBS_DISPATCH(0x34); // 7 + FAISS_QBS_DISPATCH(0x133); // 7 + FAISS_QBS_DISPATCH(0x6); // 6 + FAISS_QBS_DISPATCH(0x33); // 6 + FAISS_QBS_DISPATCH(0x123); // 6 + FAISS_QBS_DISPATCH(0x222); // 6 + FAISS_QBS_DISPATCH(0x23); // 5 + FAISS_QBS_DISPATCH(0x5); // 5 + FAISS_QBS_DISPATCH(0x13); // 4 + FAISS_QBS_DISPATCH(0x22); // 4 + FAISS_QBS_DISPATCH(0x4); // 4 + FAISS_QBS_DISPATCH(0x3); // 3 + FAISS_QBS_DISPATCH(0x21); // 3 + FAISS_QBS_DISPATCH(0x2); // 2 + FAISS_QBS_DISPATCH(0x1); // 1 +#undef FAISS_QBS_DISPATCH } - // Default: qbs not known at compile time for_each_block<32>(ntotal2, codes, block_stride, res, [&](size_t j0) { const uint8_t* LUT = LUT0; int qi = qbs; @@ -215,16 +210,16 @@ void pq4_accumulate_loop_qbs_fixed_scaler_256( int nq = qi & 15; qi >>= 4; res.set_block_origin(i0, j0); -#define FAISS_NQ256_DISPATCH(NQ) \ - case NQ: \ - pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); \ +#define FAISS_NQ_DISPATCH(NQ) \ + case NQ: \ + PQ4QBSKernel::template run(nsq, codes, LUT, res, scaler); \ break switch (nq) { - FAISS_NQ256_DISPATCH(1); - FAISS_NQ256_DISPATCH(2); - FAISS_NQ256_DISPATCH(3); - FAISS_NQ256_DISPATCH(4); -#undef FAISS_NQ256_DISPATCH + FAISS_NQ_DISPATCH(1); + FAISS_NQ_DISPATCH(2); + FAISS_NQ_DISPATCH(3); + FAISS_NQ_DISPATCH(4); +#undef FAISS_NQ_DISPATCH default: FAISS_THROW_FMT("accumulate nq=%d not instantiated", nq); } diff --git a/faiss/impl/fast_scan/accumulate_loops_512.h b/faiss/impl/fast_scan/accumulate_loops_512.h index 1b9d73b60..f756e2480 100644 --- a/faiss/impl/fast_scan/accumulate_loops_512.h +++ b/faiss/impl/fast_scan/accumulate_loops_512.h @@ -1,5 +1,13 @@ /* - * Copyright (c) Meta Platforms, Inc. and affiliates. + * Portions 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. + */ + +/* + * Portions Copyright 2026 Arm Limited and/or its affiliates + * * * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. @@ -171,11 +179,11 @@ void pq4_accumulate_loop_qbs_fixed_scaler_512( // This is rare — pq4_preferred_qbs() covers all values above. if constexpr (Scaler::nscale == 0) { DummyScaler scaler_avx2; - pq4_accumulate_loop_qbs_fixed_scaler_256( + pq4_accumulate_loop_qbs_fixed_scaler_simd( qbs, ntotal2, nsq, codes, LUT0, res, scaler_avx2, block_stride); } else { NormTableScaler scaler_avx2(scaler.scale_int); - pq4_accumulate_loop_qbs_fixed_scaler_256( + pq4_accumulate_loop_qbs_fixed_scaler_simd( qbs, ntotal2, nsq, codes, LUT0, res, scaler_avx2, block_stride); } } diff --git a/faiss/impl/fast_scan/decompose_qbs.h b/faiss/impl/fast_scan/decompose_qbs.h index dc44c3d60..9a06cd357 100644 --- a/faiss/impl/fast_scan/decompose_qbs.h +++ b/faiss/impl/fast_scan/decompose_qbs.h @@ -1,5 +1,13 @@ /* - * Copyright (c) Meta Platforms, Inc. and affiliates. + * Portions 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. + */ + +/* + * Portions Copyright 2026 Arm Limited and/or its affiliates + * * * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. @@ -12,6 +20,9 @@ #include #include #include +#ifdef COMPILE_SIMD_ARM_NEON +#include +#endif #include namespace faiss { @@ -44,59 +55,13 @@ void kernel_accumulate_block( KernelSL == SIMDLevel::AVX512_SPR) { pq4_kernel_qbs_512(nsq, codes, LUT, res, scaler); } else { - pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); + PQ4QBSKernel::template run(nsq, codes, LUT, res, scaler); } #else - pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); + PQ4QBSKernel::template run(nsq, codes, LUT, res, scaler); #endif } -// handle at most 4 blocks of queries -template < - int QBS, - SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, - class ResultHandler, - class Scaler> -void accumulate_q_4step( - size_t ntotal2, - int nsq, - const uint8_t* codes, - const uint8_t* LUT0, - ResultHandler& res, - const Scaler& scaler, - size_t block_stride) { - constexpr int Q1 = QBS & 15; - constexpr int Q2 = (QBS >> 4) & 15; - constexpr int Q3 = (QBS >> 8) & 15; - constexpr int Q4 = (QBS >> 12) & 15; - constexpr int SQ = Q1 + Q2 + Q3 + Q4; - - for_each_block<32>(ntotal2, codes, block_stride, res, [&](size_t) { - FixedStorageHandler res2; - const uint8_t* LUT = LUT0; - kernel_accumulate_block(nsq, codes, LUT, res2, scaler); - LUT += Q1 * nsq * 16; - if (Q2 > 0) { - res2.set_block_origin(Q1, 0); - kernel_accumulate_block( - nsq, codes, LUT, res2, scaler); - LUT += Q2 * nsq * 16; - } - if (Q3 > 0) { - res2.set_block_origin(Q1 + Q2, 0); - kernel_accumulate_block( - nsq, codes, LUT, res2, scaler); - LUT += Q3 * nsq * 16; - } - if (Q4 > 0) { - res2.set_block_origin(Q1 + Q2 + Q3, 0); - kernel_accumulate_block( - nsq, codes, LUT, res2, scaler); - } - res2.to_other_handler(res); - }); -} - template < int NQ, SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, @@ -149,82 +114,4 @@ void accumulate( #undef DISPATCH } -template < - SIMDLevel KernelSL = SINGLE_SIMD_LEVEL, - class ResultHandler, - class Scaler> -void pq4_accumulate_loop_qbs_fixed_scaler( - int qbs, - size_t ntotal2, - int nsq, - const uint8_t* codes, - const uint8_t* LUT0, - ResultHandler& res, - const Scaler& scaler, - size_t block_stride = 0) { - assert(nsq % 2 == 0); - assert(is_aligned_pointer(codes)); - assert(is_aligned_pointer(LUT0)); - - // try out optimized versions - switch (qbs) { -#define DISPATCH(QBS) \ - case QBS: \ - accumulate_q_4step( \ - ntotal2, nsq, codes, LUT0, res, scaler, block_stride); \ - return; - DISPATCH(0x3333); // 12 - DISPATCH(0x2333); // 11 - DISPATCH(0x2233); // 10 - DISPATCH(0x333); // 9 - DISPATCH(0x2223); // 9 - DISPATCH(0x233); // 8 - DISPATCH(0x1223); // 8 - DISPATCH(0x223); // 7 - DISPATCH(0x34); // 7 - DISPATCH(0x133); // 7 - DISPATCH(0x6); // 6 - DISPATCH(0x33); // 6 - DISPATCH(0x123); // 6 - DISPATCH(0x222); // 6 - DISPATCH(0x23); // 5 - DISPATCH(0x5); // 5 - DISPATCH(0x13); // 4 - DISPATCH(0x22); // 4 - DISPATCH(0x4); // 4 - DISPATCH(0x3); // 3 - DISPATCH(0x21); // 3 - DISPATCH(0x2); // 2 - DISPATCH(0x1); // 1 -#undef DISPATCH - } - - // default implementation where qbs is not known at compile time - for_each_block<32>(ntotal2, codes, block_stride, res, [&](size_t j0) { - const uint8_t* LUT = LUT0; - int qi = qbs; - int i0 = 0; - while (qi) { - int nq = qi & 15; - qi >>= 4; - res.set_block_origin(i0, j0); -#define DISPATCH(NQ) \ - case NQ: \ - kernel_accumulate_block(nsq, codes, LUT, res, scaler); \ - break - switch (nq) { - DISPATCH(1); - DISPATCH(2); - DISPATCH(3); - DISPATCH(4); -#undef DISPATCH - default: - FAISS_THROW_FMT("accumulate nq=%d not instantiated", nq); - } - i0 += nq; - LUT += nq * nsq * 16; - } - }); -} - } // namespace faiss diff --git a/faiss/impl/fast_scan/dispatching.h b/faiss/impl/fast_scan/dispatching.h index 2dfe5da07..731242dd6 100644 --- a/faiss/impl/fast_scan/dispatching.h +++ b/faiss/impl/fast_scan/dispatching.h @@ -1,5 +1,13 @@ /* - * Copyright (c) Meta Platforms, Inc. and affiliates. + * Portions 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. + */ + +/* + * Portions Copyright 2026 Arm Limited and/or its affiliates + * * * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. @@ -20,8 +28,8 @@ * #define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2 * #include * - * Kernel helpers come from accumulate_loops.h (search_1 multi-BB path - * and QBS 256-bit path) and accumulate_loops_512.h (QBS 512-bit path, + * Kernel helpers come from accumulate_loops.h (search_1 multi-BB path and + * SIMD-level-selected QBS path) and accumulate_loops_512.h (QBS 512-bit path, * AVX512 TU only). */ @@ -126,7 +134,7 @@ struct ScannerMixIn : FastScanCodeScanner { if (SIMDConfig::avx512_split) { if (pq2x4_scale) { NormTableScaler scaler(pq2x4_scale); - pq4_accumulate_loop_qbs_fixed_scaler_256( + pq4_accumulate_loop_qbs_fixed_scaler_simd( qbs, nb, nsq, @@ -137,7 +145,7 @@ struct ScannerMixIn : FastScanCodeScanner { block_stride); } else { DummyScaler dummy; - pq4_accumulate_loop_qbs_fixed_scaler_256( + pq4_accumulate_loop_qbs_fixed_scaler_simd( qbs, nb, nsq, @@ -174,7 +182,8 @@ struct ScannerMixIn : FastScanCodeScanner { } else { if (pq2x4_scale) { NormTableScaler scaler(pq2x4_scale); - pq4_accumulate_loop_qbs_fixed_scaler_256( + pq4_accumulate_loop_qbs_fixed_scaler_simd< + THE_LEVEL_TO_DISPATCH>( qbs, nb, nsq, @@ -185,7 +194,8 @@ struct ScannerMixIn : FastScanCodeScanner { block_stride); } else { DummyScaler dummy; - pq4_accumulate_loop_qbs_fixed_scaler_256( + pq4_accumulate_loop_qbs_fixed_scaler_simd< + THE_LEVEL_TO_DISPATCH>( qbs, nb, nsq, diff --git a/faiss/impl/fast_scan/kernels_neon.h b/faiss/impl/fast_scan/kernels_neon.h new file mode 100644 index 000000000..e3e8ecd7a --- /dev/null +++ b/faiss/impl/fast_scan/kernels_neon.h @@ -0,0 +1,155 @@ +/* + * Portions 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. + */ + +/* + * Portions Copyright 2026 Arm Limited and/or its 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 + +#include + +#include +#include +#include + +namespace faiss { + +/* + * Single-BB QBS variant: accumulates NQ queries x 32 db elements (BB=1). + * Uses native 128-bit AdvSIMD registers instead of logical 256-bit values. + */ +template +FAISS_ALWAYS_INLINE void pq4_kernel_qbs_neon( + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + ResultHandler& res, + const Scaler& scaler) { + // Avoid a zero-length array for NQ=0 instantiations. + constexpr int NQA = NQ > 0 ? NQ : 1; + const uint16x8_t zero = vdupq_n_u16(0); + // accu[q][b] holds distances for vectors 8*b..8*b+7. + uint16x8_t accu[NQA][4]; + + for (int q = 0; q < NQ; q++) { + for (int b = 0; b < 4; b++) { + accu[q][b] = zero; + } + } + + // Process two subquantizers per iteration. + for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) { + const uint8x16_t code0 = vld1q_u8(codes); + const uint8x16_t code1 = vld1q_u8(codes + 16); + codes += 32; + + const uint8x16_t mask = vdupq_n_u8(0x0f); + const uint8x16_t code0_lo = vandq_u8(code0, mask); + const uint8x16_t code1_lo = vandq_u8(code1, mask); + const uint8x16_t code0_hi = vshrq_n_u8(code0, 4); + const uint8x16_t code1_hi = vshrq_n_u8(code1, 4); + + for (int q = 0; q < NQ; q++) { + // Load LUTs for two subquantizers. + const uint8x16_t lut0 = vld1q_u8(LUT); + const uint8x16_t lut1 = vld1q_u8(LUT + 16); + LUT += 32; + + const uint8x16_t d0_lo = vqtbl1q_u8(lut0, code0_lo); + const uint8x16_t d0_hi = vqtbl1q_u8(lut0, code0_hi); + const uint8x16_t d1_lo = vqtbl1q_u8(lut1, code1_lo); + const uint8x16_t d1_hi = vqtbl1q_u8(lut1, code1_hi); + + accu[q][0] = vaddq_u16( + accu[q][0], + vaddl_u8(vget_low_u8(d0_lo), vget_low_u8(d1_lo))); + accu[q][1] = vaddq_u16(accu[q][1], vaddl_high_u8(d0_lo, d1_lo)); + accu[q][2] = vaddq_u16( + accu[q][2], + vaddl_u8(vget_low_u8(d0_hi), vget_low_u8(d1_hi))); + accu[q][3] = vaddq_u16(accu[q][3], vaddl_high_u8(d0_hi, d1_hi)); + } + } + + // Scaled tail used by additive-quantizer FastScan indexes. + // (this is compiled out for DummyScaler) + for (int sq = 0; sq < scaler.nscale; sq += 2) { + const uint16_t scale = scaler.scale_one(uint16_t{1}); + const uint8x16_t code0 = vld1q_u8(codes); + const uint8x16_t code1 = vld1q_u8(codes + 16); + codes += 32; + + const uint8x16_t mask = vdupq_n_u8(0x0f); + const uint8x16_t code0_lo = vandq_u8(code0, mask); + const uint8x16_t code1_lo = vandq_u8(code1, mask); + const uint8x16_t code0_hi = vshrq_n_u8(code0, 4); + const uint8x16_t code1_hi = vshrq_n_u8(code1, 4); + + for (int q = 0; q < NQ; q++) { + // Load LUTs for two subquantizers. + const uint8x16_t lut0 = vld1q_u8(LUT); + const uint8x16_t lut1 = vld1q_u8(LUT + 16); + LUT += 32; + + const uint8x16_t d0_lo = vqtbl1q_u8(lut0, code0_lo); + const uint8x16_t d0_hi = vqtbl1q_u8(lut0, code0_hi); + const uint8x16_t d1_lo = vqtbl1q_u8(lut1, code1_lo); + const uint8x16_t d1_hi = vqtbl1q_u8(lut1, code1_hi); + + accu[q][0] = vaddq_u16( + accu[q][0], + vmulq_n_u16( + vaddl_u8(vget_low_u8(d0_lo), vget_low_u8(d1_lo)), + scale)); + accu[q][1] = vaddq_u16( + accu[q][1], + vmulq_n_u16(vaddl_high_u8(d0_lo, d1_lo), scale)); + accu[q][2] = vaddq_u16( + accu[q][2], + vmulq_n_u16( + vaddl_u8(vget_low_u8(d0_hi), vget_low_u8(d1_hi)), + scale)); + accu[q][3] = vaddq_u16( + accu[q][3], + vmulq_n_u16(vaddl_high_u8(d0_hi, d1_hi), scale)); + } + } + + // Reorder the four 8-lane accumulators for the result handler. + using simd16uint16 = simd16uint16_tpl; + for (int q = 0; q < NQ; q++) { + const uint16x8_t d0 = vuzp1q_u16(accu[q][0], accu[q][1]); + const uint16x8_t d1 = vuzp2q_u16(accu[q][0], accu[q][1]); + const uint16x8_t d2 = vuzp1q_u16(accu[q][2], accu[q][3]); + const uint16x8_t d3 = vuzp2q_u16(accu[q][2], accu[q][3]); + res.handle( + q, + 0, + simd16uint16(uint16x8x2_t{d0, d1}), + simd16uint16(uint16x8x2_t{d2, d3})); + } +} + +template <> +struct PQ4QBSKernel { + template + FAISS_ALWAYS_INLINE static void run( + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + ResultHandler& res, + const Scaler& scaler) { + pq4_kernel_qbs_neon(nsq, codes, LUT, res, scaler); + } +}; + +} // namespace faiss diff --git a/faiss/impl/fast_scan/kernels_simd256.h b/faiss/impl/fast_scan/kernels_simd256.h index 71142832e..721b9e871 100644 --- a/faiss/impl/fast_scan/kernels_simd256.h +++ b/faiss/impl/fast_scan/kernels_simd256.h @@ -1,5 +1,13 @@ /* - * Copyright (c) Meta Platforms, Inc. and affiliates. + * Portions 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. + */ + +/* + * Portions Copyright 2026 Arm Limited and/or its affiliates + * * * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. @@ -210,4 +218,17 @@ void pq4_kernel_qbs_256( } } +template +struct PQ4QBSKernel { + template + FAISS_ALWAYS_INLINE static void run( + int nsq, + const uint8_t* codes, + const uint8_t* LUT, + ResultHandler& res, + const Scaler& scaler) { + pq4_kernel_qbs_256(nsq, codes, LUT, res, scaler); + } +}; + } // namespace faiss diff --git a/tests/test_single_query_collect_handler.cpp b/tests/test_single_query_collect_handler.cpp index ad3e05d85..ea2cfb444 100644 --- a/tests/test_single_query_collect_handler.cpp +++ b/tests/test_single_query_collect_handler.cpp @@ -1,5 +1,13 @@ /* - * Copyright (c) Meta Platforms, Inc. and affiliates. + * Portions 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. + */ + +/* + * Portions Copyright 2026 Arm Limited and/or its affiliates + * * * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. @@ -10,7 +18,7 @@ * * Drives the handler through the actual fast_scan accumulate loop * by building an IndexPQFastScan, computing its quantized LUT, - * and calling pq4_accumulate_loop_qbs_fixed_scaler_256 directly. + * and calling pq4_accumulate_loop_qbs_fixed_scaler_simd directly. */ #include @@ -89,7 +97,7 @@ std::vector> run_collect_handler( // 5. Run the accumulate loop DummyScaler<> scaler; - pq4_accumulate_loop_qbs_fixed_scaler_256( + pq4_accumulate_loop_qbs_fixed_scaler_simd( qbs, index.ntotal2, index.M2,