feat: AMX-BF16 inner-product distance computer for IndexHNSWSQ

[xtangxtang]

Summary

Accelerate IndexHNSWSQ(QT_bf16) inner-product search on Intel AMX
(Sapphire Rapids / Granite Rapids and newer) by computing 16 bf16 inner
products per distances_batch_16 call in a single AMX tile pass.

Re-architected from the original standalone-storage version to fit faiss's
existing abstractions (per review):

1. BF16 via the ScalarQuantizer (IndexHNSWSQ)

QuantizerBF16 (QT_bf16) already stores each component as bf16, so the AMX
tile engine consumes the SQ codes directly — no separate bf16_storage, no
re-encoding, and no bf16-specific members on IndexHNSW.

2. Batched DistanceComputer, 4 → 16

DistanceComputer::distances_batch_16 (default = 4× distances_batch_4, so
every computer benefits for free) is overridden by the AMX bf16 computer. The
level-0 HNSW loop (search_from_candidates_fixVT) now batches 16 neighbours;
the upper-level greedy and unbounded loops keep batch-4.

3. AMX behind the SIMD dynamic dispatch

New SIMDLevel::AMX (above AVX512_SPR) with runtime CPUID + tile-permission
detection and fallback chain AMX → AVX512_SPR → AVX512 → AVX2 → NONE. The
kernel is faiss/impl/scalar_quantizer/sq-amx.cpp (guarded by
COMPILE_SIMD_AMX, compiled into a faiss_amx target with
-mamx-tile -mamx-bf16, mirroring faiss_avx512_spr).
ScalarQuantizer::get_distance_computer dispatches via with_simd_level_amx.

Scope: inner product only

AMX is routed only for QT_bf16 + METRIC_INNER_PRODUCT. L2 (and every other
quantizer type) uses the AVX512 SQ path, because the per-vector ‖c‖² needed
to turn the tile dot product into an L2 distance offsets the tile gain.

Performance

Granite Rapids (Xeon 6972P), d=768, 100k vectors, efSearch=128; AMX vs the
AVX512 bf16 SQ path on an identical index, recall unchanged (~0.99):

threads	AVX512 QPS	AMX QPS	speedup
1	1,386	1,831	1.32×
16 (8C/16T)	15,008	18,217	1.21×

At full-socket thread counts the workload becomes memory-bandwidth-bound and
the two converge; the gain is in the compute-bound regime.

Usage

faiss::IndexHNSWSQ index(
    d, faiss::ScalarQuantizer::QT_bf16, M, faiss::METRIC_INNER_PRODUCT);
index.train(n, xb);
index.add(n, xb);
index.search(nq, xq, k, distances, labels);  // AMX tile path on SPR/GNR+

Commits

1. feat: add AMX SIMD level to the dynamic dispatch framework
2. feat: add DistanceComputer::distances_batch_16 for HNSW level-0 search
3. feat: add AMX-BF16 tile inner-product distance computer

[mdouze]

Thanks for the work on this PR.

In order for us to accept it, it should be re-architctured in the following way:

• BF16 is supported in the ScalarQuantizer, see https://github.com/facebookresearch/faiss/blob/main/faiss/impl/ScalarQuantizer.h#L33, so the encoding implementation should go in the IndexHNSWSQ.

• the DistanceComputer interface already supports batches of 4, see here https://github.com/facebookresearch/faiss/blob/main/faiss/impl/HNSW.cpp#L976, it should be extended to batches of 16 and specialized for AMX

• the AMX code should be inside a specific SIMD file and enabled via the dispatching mechanism described here https://github.com/facebookresearch/faiss/blob/main/faiss/docs/simd_dynamic_dispatch_migration.md

Thanks!

[xtangxtang]

Thanks for the detailed review, @mdouze — I've re-architected the PR along all three points. The standalone BF16 storage + AMX path is gone; everything now flows through the existing ScalarQuantizer and SIMD dynamic-dispatch infrastructure.

1. BF16 encoding lives in the ScalarQuantizer / IndexHNSWSQ
The separate bf16_storage / bf16_norms members on IndexHNSW are removed. Storage and query encoding now use IndexHNSWSQ(d, ScalarQuantizer::QT_bf16, M). Since QuantizerBF16 already stores each component as bf16, the AMX tile engine consumes the SQ codes directly — no separate storage and no re-encoding.

2. The batched DistanceComputer interface, extended 4 → 16
Added DistanceComputer::distances_batch_16, with a default implementation that groups the work into 4× distances_batch_4 (so every existing computer benefits for free), overridden by the AMX bf16 computer. The level-0 HNSW loop (search_from_candidates_fixVT) now accumulates 16 neighbours and calls distances_batch_16; the upper-level greedy and unbounded loops keep batch-4.

3. AMX in a dedicated SIMD file, behind the dispatch mechanism
Added SIMDLevel::AMX (above AVX512_SPR) to simd_levels.{h,cpp} and simd_dispatch.h — runtime CPUID + tile-permission detection, fallback chain AMX → AVX512_SPR → AVX512 → AVX2 → NONE. The kernel is faiss/impl/scalar_quantizer/sq-amx.cpp, guarded by COMPILE_SIMD_AMX and built into a faiss_amx target with -mamx-tile -mamx-bf16 (mirroring faiss_avx512_spr). ScalarQuantizer::get_distance_computer dispatches through with_simd_level_amx.

Scope — inner product only. AMX is routed only for QT_bf16 + METRIC_INNER_PRODUCT. For L2, the per-vector ‖c‖² needed to turn the tile dot product into a distance offsets the tile gain (and slightly hurts recall by quantising the query), so L2 and every other quantizer type use the AVX512 SQ path.

Performance (Granite Rapids, Xeon 6972P; d=768, 100k vectors, efSearch=128; AMX vs the AVX512 bf16 SQ path on an identical index, recall unchanged at ~0.99):

threads	AVX512 QPS	AMX QPS	speedup
1	1,386	1,831	1.32×
16 (8C/16T)	15,008	18,217	1.21×

At full-socket thread counts the workload becomes memory-bandwidth-bound and the two converge, as expected — the gain is in the compute-bound regime (latency / moderate concurrency).

The branch is rebased onto current main and is now three focused commits (AMX SIMD level / distances_batch_16 / sq-amx.cpp). I'll update the PR title and description to match the new scope.