mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
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
This commit is contained in:
Fadi Arafeh
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
7069478c3f
commit
a2f7943005
@@ -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 <[email protected]>
|
||||
#
|
||||
# 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
|
||||
|
||||
@@ -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
|
||||
* <[email protected]>
|
||||
*
|
||||
* 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<SL> 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 <cassert>
|
||||
@@ -37,6 +37,9 @@
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/fast_scan/LookupTableScaler.h>
|
||||
#include <faiss/impl/fast_scan/kernels_simd256.h>
|
||||
#ifdef COMPILE_SIMD_ARM_NEON
|
||||
#include <faiss/impl/fast_scan/kernels_neon.h>
|
||||
#endif
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
|
||||
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 <int QBS, SIMDLevel SL, class ResultHandler, class Scaler>
|
||||
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<SQ, 2, KernelSL> res2;
|
||||
FixedStorageHandler<SQ, 2, SL> res2;
|
||||
const uint8_t* LUT = LUT0;
|
||||
pq4_kernel_qbs_256<Q1, KernelSL>(nsq, codes, LUT, res2, scaler);
|
||||
PQ4QBSKernel<SL>::template run<Q1>(nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q1 * nsq * 16;
|
||||
if (Q2 > 0) {
|
||||
res2.set_block_origin(Q1, 0);
|
||||
pq4_kernel_qbs_256<Q2, KernelSL>(nsq, codes, LUT, res2, scaler);
|
||||
PQ4QBSKernel<SL>::template run<Q2>(nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q2 * nsq * 16;
|
||||
}
|
||||
if (Q3 > 0) {
|
||||
res2.set_block_origin(Q1 + Q2, 0);
|
||||
pq4_kernel_qbs_256<Q3, KernelSL>(nsq, codes, LUT, res2, scaler);
|
||||
PQ4QBSKernel<SL>::template run<Q3>(nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q3 * nsq * 16;
|
||||
}
|
||||
if (Q4 > 0) {
|
||||
res2.set_block_origin(Q1 + Q2 + Q3, 0);
|
||||
pq4_kernel_qbs_256<Q4, KernelSL>(nsq, codes, LUT, res2, scaler);
|
||||
PQ4QBSKernel<SL>::template run<Q4>(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 <SIMDLevel SL = SINGLE_SIMD_LEVEL, class ResultHandler, class Scaler>
|
||||
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<QBS, KernelSL>( \
|
||||
accumulate_q_4step_for_simd<QBS, SL>( \
|
||||
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<NQ, KernelSL>(nsq, codes, LUT, res, scaler); \
|
||||
#define FAISS_NQ_DISPATCH(NQ) \
|
||||
case NQ: \
|
||||
PQ4QBSKernel<SL>::template run<NQ>(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);
|
||||
}
|
||||
|
||||
@@ -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
|
||||
* <[email protected]>
|
||||
*
|
||||
* 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<SIMDLevel::AVX2> scaler_avx2;
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256<SIMDLevel::AVX2>(
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_simd<SIMDLevel::AVX2>(
|
||||
qbs, ntotal2, nsq, codes, LUT0, res, scaler_avx2, block_stride);
|
||||
} else {
|
||||
NormTableScaler<SIMDLevel::AVX2> scaler_avx2(scaler.scale_int);
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256<SIMDLevel::AVX2>(
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_simd<SIMDLevel::AVX2>(
|
||||
qbs, ntotal2, nsq, codes, LUT0, res, scaler_avx2, block_stride);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
* <[email protected]>
|
||||
*
|
||||
* 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 <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/fast_scan/kernels_simd256.h>
|
||||
#include <faiss/impl/fast_scan/kernels_simd512.h>
|
||||
#ifdef COMPILE_SIMD_ARM_NEON
|
||||
#include <faiss/impl/fast_scan/kernels_neon.h>
|
||||
#endif
|
||||
#include <faiss/impl/fast_scan/simd_result_handlers.h>
|
||||
|
||||
namespace faiss {
|
||||
@@ -44,59 +55,13 @@ void kernel_accumulate_block(
|
||||
KernelSL == SIMDLevel::AVX512_SPR) {
|
||||
pq4_kernel_qbs_512<NQ>(nsq, codes, LUT, res, scaler);
|
||||
} else {
|
||||
pq4_kernel_qbs_256<NQ, KernelSL>(nsq, codes, LUT, res, scaler);
|
||||
PQ4QBSKernel<KernelSL>::template run<NQ>(nsq, codes, LUT, res, scaler);
|
||||
}
|
||||
#else
|
||||
pq4_kernel_qbs_256<NQ, KernelSL>(nsq, codes, LUT, res, scaler);
|
||||
PQ4QBSKernel<KernelSL>::template run<NQ>(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<SQ, 2, KernelSL> res2;
|
||||
const uint8_t* LUT = LUT0;
|
||||
kernel_accumulate_block<Q1, KernelSL>(nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q1 * nsq * 16;
|
||||
if (Q2 > 0) {
|
||||
res2.set_block_origin(Q1, 0);
|
||||
kernel_accumulate_block<Q2, KernelSL>(
|
||||
nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q2 * nsq * 16;
|
||||
}
|
||||
if (Q3 > 0) {
|
||||
res2.set_block_origin(Q1 + Q2, 0);
|
||||
kernel_accumulate_block<Q3, KernelSL>(
|
||||
nsq, codes, LUT, res2, scaler);
|
||||
LUT += Q3 * nsq * 16;
|
||||
}
|
||||
if (Q4 > 0) {
|
||||
res2.set_block_origin(Q1 + Q2 + Q3, 0);
|
||||
kernel_accumulate_block<Q4, KernelSL>(
|
||||
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<QBS, KernelSL>( \
|
||||
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<NQ, KernelSL>(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
|
||||
|
||||
@@ -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
|
||||
* <[email protected]>
|
||||
*
|
||||
* 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 <faiss/impl/fast_scan/dispatching.h>
|
||||
*
|
||||
* 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<SIMDLevel::AVX2> scaler(pq2x4_scale);
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256<SIMDLevel::AVX2>(
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_simd<SIMDLevel::AVX2>(
|
||||
qbs,
|
||||
nb,
|
||||
nsq,
|
||||
@@ -137,7 +145,7 @@ struct ScannerMixIn : FastScanCodeScanner {
|
||||
block_stride);
|
||||
} else {
|
||||
DummyScaler<SIMDLevel::AVX2> dummy;
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256<SIMDLevel::AVX2>(
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_simd<SIMDLevel::AVX2>(
|
||||
qbs,
|
||||
nb,
|
||||
nsq,
|
||||
@@ -174,7 +182,8 @@ struct ScannerMixIn : FastScanCodeScanner {
|
||||
} else {
|
||||
if (pq2x4_scale) {
|
||||
NormTableScaler<THE_LEVEL_TO_DISPATCH> scaler(pq2x4_scale);
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256<THE_LEVEL_TO_DISPATCH>(
|
||||
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<THE_LEVEL_TO_DISPATCH> dummy;
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_256<THE_LEVEL_TO_DISPATCH>(
|
||||
pq4_accumulate_loop_qbs_fixed_scaler_simd<
|
||||
THE_LEVEL_TO_DISPATCH>(
|
||||
qbs,
|
||||
nb,
|
||||
nsq,
|
||||
|
||||
@@ -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
|
||||
* <[email protected]>
|
||||
*
|
||||
* 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 <arm_neon.h>
|
||||
|
||||
#include <faiss/impl/fast_scan/kernels_simd256.h>
|
||||
#include <faiss/impl/platform_macros.h>
|
||||
#include <faiss/impl/simdlib/simdlib_dispatch.h>
|
||||
|
||||
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 <int NQ, class ResultHandler, class Scaler>
|
||||
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<SIMDLevel::ARM_NEON>;
|
||||
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<SIMDLevel::ARM_NEON> {
|
||||
template <int NQ, class ResultHandler, class Scaler>
|
||||
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<NQ>(nsq, codes, LUT, res, scaler);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace faiss
|
||||
@@ -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
|
||||
* <[email protected]>
|
||||
*
|
||||
* 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 <SIMDLevel SL>
|
||||
struct PQ4QBSKernel {
|
||||
template <int NQ, class ResultHandler, class Scaler>
|
||||
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<NQ, SL>(nsq, codes, LUT, res, scaler);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace faiss
|
||||
|
||||
@@ -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
|
||||
* <[email protected]>
|
||||
*
|
||||
* 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 <algorithm>
|
||||
@@ -89,7 +97,7 @@ std::vector<std::pair<int64_t, float>> 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,
|
||||
|
||||
Reference in New Issue
Block a user