Skip to content

Name the SIMD level masks after what they hold #5600

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mnorris11 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Name the SIMD level masks after what they hold #5600

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c92f7ca
Stop aarch64 falling back to SIMDLevel::NONE in faiss dispatch (#5572)
Sep 7, 2026
6193c2b
Name the SIMD level masks after what they hold
Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions faiss/IndexFlat.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,10 +269,15 @@ struct FlatIPDis : FlatCodesDistanceComputer {

FlatCodesDistanceComputer* IndexFlat::get_FlatCodesDistanceComputer() const {
FlatCodesDistanceComputer* dc = nullptr;
// BASE_WITH_SVE rather than plain with_simd_level: the fvec_* kernels
// these computers call have dedicated ARM_SVE specializations, which the
// BASE mask would skip in favour of ARM_NEON.
if (metric_type == METRIC_L2) {
with_simd_level([&]<SIMDLevel SL>() { dc = new FlatL2Dis<SL>(*this); });
with_simd_level_with_sve(
[&]<SIMDLevel SL>() { dc = new FlatL2Dis<SL>(*this); });
} else if (metric_type == METRIC_INNER_PRODUCT) {
with_simd_level([&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
with_simd_level_with_sve(
[&]<SIMDLevel SL>() { dc = new FlatIPDis<SL>(*this); });
} else {
dc = get_extra_distance_computer(d, metric_type, metric_arg, get_xb());
}
Expand Down Expand Up @@ -404,7 +409,7 @@ FlatCodesDistanceComputer* IndexFlatL2::get_FlatCodesDistanceComputer() const {
if (metric_type == METRIC_L2) {
if (!cached_l2norms.empty()) {
FlatCodesDistanceComputer* dc = nullptr;
with_simd_level([&]<SIMDLevel SL>() {
with_simd_level_with_sve([&]<SIMDLevel SL>() {
dc = new FlatL2WithNormsDis<SL>(*this);
});
return dc;
Expand Down Expand Up @@ -784,7 +789,9 @@ void IndexFlatPanorama::search_subset(
idx_t k,
float* distances,
idx_t* labels) const {
with_simd_level([&]<SIMDLevel SL>() {
// BASE_WITH_SVE: the only SIMD-templated call in the loop body is
// fvec_inner_product, which has an ARM_SVE specialization.
with_simd_level_with_sve([&]<SIMDLevel SL>() {
with_metric_type(metric_type, [&]<MetricType M>() {
constexpr bool is_sim = is_similarity_metric(M);
using C = std::conditional_t<
Expand Down
7 changes: 5 additions & 2 deletions faiss/SuperKMeans.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -620,8 +620,11 @@ void super_kmeans_assign_iteration(
&ldc);
}

// One SIMD dispatch per (xi, yj) tile.
with_simd_level([&]<SIMDLevel SL>() {
// One SIMD dispatch per (xi, yj) tile. BASE_WITH_SVE rather than
// plain with_simd_level: block_l2 has a dedicated ARM_SVE kernel,
// which the BASE mask would skip in favour of ARM_NEON (no
// specialization, so the scalar primary template).
with_simd_level_with_sve([&]<SIMDLevel SL>() {
[[maybe_unused]] const int omp_chunk_local = cp.omp_chunk;
int64_t tile_total = 0;
int64_t tile_pruned = 0;
Expand Down
5 changes: 3 additions & 2 deletions faiss/docs/simd_dynamic_dispatch_migration.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,8 +238,9 @@ your own with `(1 << int(SIMDLevel::X)) | ...`):
|------|--------|---------|
| `AVAILABLE_SIMD_LEVELS_NONE` | NONE only | Scalar-only functions |
| `AVAILABLE_SIMD_LEVELS_AVX2_NEON` | NONE, AVX2, ARM_NEON | 256-bit `simdlib` ops (`with_simd_level_256bit`) |
| `AVAILABLE_SIMD_LEVELS_A0` | NONE, AVX2, AVX512, ARM_NEON, RISCV_RVV | Default (`with_simd_level`) |
| `AVAILABLE_SIMD_LEVELS_A1` | A0 + ARM_SVE | Functions with dedicated SVE implementations |
| `AVAILABLE_SIMD_LEVELS_BASE` | NONE, AVX2, AVX512, ARM_NEON, RISCV_RVV | Default (`with_simd_level`). ARM_NEON is part of BASE: NEON is mandatory on aarch64, SVE is optional |
| `AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR` | BASE + AVX512_SPR | Kernels needing the whole SPR feature set (`with_simd_level_with_spr`) |
| `AVAILABLE_SIMD_LEVELS_BASE_WITH_SVE` | BASE + ARM_SVE | Functions with dedicated SVE implementations (`with_simd_level_with_sve`) |
| `AVAILABLE_SIMD_LEVELS_ALL` | All levels | Identity / diagnostic functions |

### Step 4: Register in build system
Expand Down
3 changes: 2 additions & 1 deletion faiss/impl/AdditiveQuantizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,7 +376,8 @@ void AdditiveQuantizer::compute_centroid_norms(float* norms) const {
size_t ntotal = (size_t)1 << tot_bits;
int64_t ntotal_signed = ntotal;
// TODO: make tree of partial sums
with_simd_level([&]<SIMDLevel SL>() {
// BASE_WITH_SVE: fvec_norm_L2sqr has an ARM_SVE specialization.
with_simd_level_with_sve([&]<SIMDLevel SL>() {
#pragma omp parallel
{
std::vector<float> tmp(d);
Expand Down
15 changes: 8 additions & 7 deletions faiss/impl/ProductQuantizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,8 +280,9 @@ void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
} // namespace

void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
// a1: fvec_L2sqr_ny_nearest / _y_transposed have ARM_SVE specializations
with_simd_level_a1([&]<SIMDLevel SL>() {
// with_sve: fvec_L2sqr_ny_nearest / _y_transposed have ARM_SVE
// specializations
with_simd_level_with_sve([&]<SIMDLevel SL>() {
switch (nbits) {
case 8:
compute_1_code<PQEncoder8, SL>(*this, x, code);
Expand All @@ -295,7 +296,7 @@ void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
compute_1_code<PQEncoderGeneric, SL>(*this, x, code);
break;
}
}); // with_simd_level_a1
}); // with_simd_level_with_sve
}

template <class PQDecoder>
Expand Down Expand Up @@ -443,8 +444,8 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)

void ProductQuantizer::compute_distance_table(const float* x, float* dis_table)
const {
// a1: fvec_L2sqr_ny / _transposed have ARM_SVE specializations
with_simd_level_a1([&]<SIMDLevel SL>() {
// with_sve: fvec_L2sqr_ny / _transposed have ARM_SVE specializations
with_simd_level_with_sve([&]<SIMDLevel SL>() {
if (transposed_centroids.empty()) {
// use regular version
for (size_t m = 0; m < M; m++) {
Expand Down Expand Up @@ -826,8 +827,8 @@ void ProductQuantizer::compute_sdc_table() {
sdc_table.resize(M * ksub * ksub);

if (dsub < 4) {
// a1: fvec_L2sqr_ny has an ARM_SVE specialization
with_simd_level_a1([&]<SIMDLevel SL>() {
// with_sve: fvec_L2sqr_ny has an ARM_SVE specialization
with_simd_level_with_sve([&]<SIMDLevel SL>() {
#pragma omp parallel for
for (int64_t mk = 0; mk < static_cast<int64_t>(M * ksub); mk++) {
// allow omp to schedule in a more fine-grained way
Expand Down
2 changes: 1 addition & 1 deletion faiss/impl/RaBitQUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -329,7 +329,7 @@ float compute_full_multibit_distance(
size_t d,
size_t ex_bits,
MetricType metric_type) {
return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0_SPR>(
return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR>(
[&]<SIMDLevel SL>() {
return compute_full_multibit_distance<SL>(
sign_bits,
Expand Down
4 changes: 2 additions & 2 deletions faiss/impl/RaBitQuantizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -665,12 +665,12 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
// call the SIMD-specialized rabitq functions directly (no per-call
// with_simd_level overhead).
//
// Use A0_SPR (which includes AVX512_SPR) so that on Sapphire Rapids
// Use BASE_WITH_SPR (which includes AVX512_SPR) so that on Sapphire Rapids
// and later x86 microarchitectures the VPOPCNTDQ-based RaBitQ
// specialization in rabitq_avx512_spr.cpp is selected. On AVX-512
// CPUs without VPOPCNTDQ, dispatch falls through to the AVX512
// specialization in rabitq_avx512.cpp.
return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0_SPR>(
return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR>(
[&]<SIMDLevel SL>() -> FlatCodesDistanceComputer* {
if (qb == 0) {
auto dc =
Expand Down
6 changes: 3 additions & 3 deletions faiss/impl/ScalarQuantizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -665,7 +665,7 @@ ScalarQuantizer::SQuantizer* ScalarQuantizer::select_quantizer() const {
// A SIMD level's factory returns nullptr when the dimension is
// incompatible (e.g. AVX-512 needs d % 16 == 0); the dispatcher then falls
// back to the next-lower level (AVX-512 -> AVX2 -> scalar).
return with_simd_level_fallback<AVAILABLE_SIMD_LEVELS_A0_SPR>(
return with_simd_level_fallback<AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR>(
[&]<SIMDLevel SL>() -> SQuantizer* {
return scalar_quantizer::sq_select_quantizer<SL>(
qtype, d, trained);
Expand Down Expand Up @@ -702,7 +702,7 @@ void ScalarQuantizer::decode(const uint8_t* codes, float* x, size_t n) const {
ScalarQuantizer::SQDistanceComputer* ScalarQuantizer::get_distance_computer(
MetricType metric) const {
FAISS_THROW_IF_NOT(metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT);
return with_simd_level_fallback<AVAILABLE_SIMD_LEVELS_A0_SPR>(
return with_simd_level_fallback<AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR>(
[&]<SIMDLevel SL>() -> SQDistanceComputer* {
return scalar_quantizer::sq_select_distance_computer<SL>(
metric, qtype, d, trained);
Expand All @@ -715,7 +715,7 @@ InvertedListScanner* ScalarQuantizer::select_InvertedListScanner(
bool store_pairs,
const IDSelector* sel,
bool by_residual) const {
return with_simd_level_fallback<AVAILABLE_SIMD_LEVELS_A0_SPR>(
return with_simd_level_fallback<AVAILABLE_SIMD_LEVELS_BASE_WITH_SPR>(
[&]<SIMDLevel SL>() -> InvertedListScanner* {
return scalar_quantizer::sq_select_InvertedListScanner<SL>(
qtype,
Expand Down
12 changes: 8 additions & 4 deletions faiss/impl/pq_code_distance/pq_code_distance-generic.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,17 @@

// This TU provides non-templated PQ code distance dispatch wrappers
// (pq_code_distance_8bit_single, pq_code_distance_8bit_four) declared
// in pq_code_distance-inl.h. These use with_simd_level to route to the
// in pq_code_distance-inl.h. These use with_simd_level_with_sve to route to the
// best available SIMD implementation via pq_code_distance_8bit_*_impl
// function template specializations.
//
// The NONE and ARM_NEON _impl specializations are defined inline in
// pq_code_distance-generic.h (included transitively). The AVX2, AVX512,
// and ARM_SVE specializations are in their respective per-SIMD files.
//
// BASE_WITH_SVE rather than plain with_simd_level: all three _impl
// functions have dedicated ARM_SVE specializations, which the BASE mask
// would skip in favour of ARM_NEON.

#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>

Expand All @@ -34,7 +38,7 @@ void pq_scan_8bit(
float* heap_dis,
int64_t* heap_ids,
bool max_heap) {
with_simd_level([&]<SIMDLevel SL>() {
with_simd_level_with_sve([&]<SIMDLevel SL>() {
pq_scan_8bit_impl<SL>(
M, dis_table, codes, ncodes, k, heap_dis, heap_ids, max_heap);
});
Expand All @@ -44,7 +48,7 @@ float pq_code_distance_8bit_single(
size_t M,
const float* sim_table,
const uint8_t* code) {
return with_simd_level([&]<SIMDLevel SL>() {
return with_simd_level_with_sve([&]<SIMDLevel SL>() {
return pq_code_distance_8bit_single_impl<SL>(M, sim_table, code);
});
}
Expand All @@ -60,7 +64,7 @@ void pq_code_distance_8bit_four(
float& result1,
float& result2,
float& result3) {
with_simd_level([&]<SIMDLevel SL>() {
with_simd_level_with_sve([&]<SIMDLevel SL>() {
pq_code_distance_8bit_four_impl<SL>(
M,
sim_table,
Expand Down
7 changes: 5 additions & 2 deletions faiss/impl/scalar_quantizer/distance_computers.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,8 +78,11 @@ template <class Similarity, SIMDLevel SL>
struct DistanceComputerByte : SQDistanceComputer {};

// Byte-domain distance computer for QT_8bit_direct_signed (storage is
// value+128). Only specialized for AVX512_SPR; other levels fall back to
// the float-domain DCTemplate path via the dispatch logic.
// value+128). Specialized for AVX2, AVX512, AVX512_SPR and ARM_NEON; levels
// without a specialization fall back to the float-domain DCTemplate path via
// the dispatch logic. Adding a level to the dispatch chain in sq-dispatch.h
// without a specialization here instantiates this empty primary template,
// whose pure virtuals are left unimplemented.
template <class Similarity, SIMDLevel SL>
struct DistanceComputerByteSigned : SQDistanceComputer {};

Expand Down
12 changes: 8 additions & 4 deletions faiss/impl/scalar_quantizer/sq-dispatch.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -531,7 +531,8 @@ SQDistanceComputer* select_distance_computer_body(
return new DistanceComputerByte<Sim, SL2>(
static_cast<int>(d), trained);
}
} else if constexpr (SL2 == SIMDLevel::AVX2) {
} else if constexpr (
SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
if (d % 16 == 0) {
return new DistanceComputerByte<Sim, SL2>(
static_cast<int>(d), trained);
Expand All @@ -551,7 +552,8 @@ SQDistanceComputer* select_distance_computer_body(
return new DistanceComputerByteSigned<Sim, SL2>(
static_cast<int>(d), trained);
}
} else if constexpr (SL2 == SIMDLevel::AVX2) {
} else if constexpr (
SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
if (d % 16 == 0) {
return new DistanceComputerByteSigned<Sim, SL2>(
static_cast<int>(d), trained);
Expand Down Expand Up @@ -744,7 +746,8 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
return scan.template
operator()<DistanceComputerByte<Similarity, SL2>>();
}
} else if constexpr (SL2 == SIMDLevel::AVX2) {
} else if constexpr (
SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
if (d % 16 == 0) {
return scan.template
operator()<DistanceComputerByte<Similarity, SL2>>();
Expand All @@ -765,7 +768,8 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
return scan.template operator()<
DistanceComputerByteSigned<Similarity, SL2>>();
}
} else if constexpr (SL2 == SIMDLevel::AVX2) {
} else if constexpr (
SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
if (d % 16 == 0) {
return scan.template operator()<
DistanceComputerByteSigned<Similarity, SL2>>();
Expand Down
Loading
Loading