Stop aarch64 falling back to SIMDLevel::NONE in faiss dispatch (#5572)

[mnorris11]

Summary:

TL;DR: Stops aarch64 running scalar code where a NEON or SVE kernel exists, adds the two missing NEON byte-domain kernels, and fixes FAISS_SIMD_LEVEL falling all the way to NONE on an uncompiled level.

On aarch64 several faiss dispatch paths ran scalar code, or ran an ARM_NEON kernel where an ARM_SVE kernel already existed. This change enforces the order SVE -> NEON -> NONE on an ARM host. It also adds the one kernel pair that was missing.

T287037898 reported the scalar-quantizer half. On aarch64 IndexScalarQuantizer with QT_8bit_direct or QT_8bit_direct_signed saved memory but lost throughput. Refine(SQ8) has the same problem, because IndexRefine::search calls refine_index->get_distance_computer().

What each call site gets

Call site	Before	After
QT_8bit_direct on aarch64	float-domain DCTemplate	NEON DistanceComputerByte
QT_8bit_direct_signed on aarch64	float-domain DCTemplate	NEON DistanceComputerByteSigned (new)
IndexFlat distance computers, SVE host	ARM_NEON	ARM_SVE
AdditiveQuantizer::compute_centroid_norms, SVE host	ARM_NEON	ARM_SVE
SuperKMeans block_l2, SVE host	ARM_NEON	ARM_SVE
pq_code_distance wrappers, SVE host	ARM_NEON	ARM_SVE
with_VectorDistance, SVE host	ARM_NEON	ARM_SVE
FAISS_SIMD_LEVEL=ARM_SVE, build without SVE	NONE	nearest compiled level

Three causes, fixed separately

1. Two missing kernels. sq-neon.cpp held a scalar DistanceComputerByte<Sim, ARM_NEON> that nothing could reach, and no DistanceComputerByteSigned<Sim, ARM_NEON> at all. Both are now real NEON kernels. L2 and the unsigned inner product use vabdq_u8, vmull_u8 and vpadalq_u16. The bias-encoded inner product uses veorq_u8, vmull_s8 and vpadalq_s16. Both agree bit for bit with the AVX2 specializations, which the tests require.

Two invariants deserve a note:

• The L2 path keeps an unsigned accumulator. A vmull_u8 square reaches 255^2 = 65025, which an int16 lane would read as negative.
• For x in 0 to 255, x ^ 0x80 read as int8 is exactly x - 128. That is how the kernel removes the +128 bias before vmull_s8.

2. Dispatch chains that list only x86 levels. The if constexpr chains in sq-dispatch.h enumerated x86 levels only, so an ARM host fell through to the float path. This adds ARM_NEON to four chains: two in select_distance_computer_body, and two in sq_select_InvertedListScanner. The chain in is_dimension_compatible already included ARM.

This does not add ARM_SVE. The scalar-quantizer entry points dispatch with a mask that holds no ARM_SVE bit, so an SVE host already falls through to the ARM_NEON case and now gets these kernels. Adding ARM_SVE would instantiate the empty primary template in distance_computers.h. This corrects that template's stale comment.

3. Level masks that hide existing SVE kernels. The default mask holds no ARM_SVE bit, so the dispatch fell through case ARM_SVE to ARM_NEON. This uses with_simd_level_a1 at every site where a real SVE kernel exists and links: IndexFlat.cpp at four sites, AdditiveQuantizer::compute_centroid_norms, block_l2 in SuperKMeans.cpp, all three wrappers in pq_code_distance-generic.cpp, and with_VectorDistance in distances_dispatch.h.

The IndexFlat sites matter most. faiss::fvec_L2sqr() already used the SVE mask, so on an SVE host the free function ran SVE while IndexFlatL2's distance computer ran NEON.

Also: FAISS_SIMD_LEVEL=ARM_SVE on a build without SVE compiled in used to skip every level and run at NONE, because with_selected_simd_levels has no case label for an uncompiled level. It now walks down to the nearest compiled level. The override is still honoured when the CPU lacks the level, since forcing a level is the point of it. Only uncompiled levels are corrected.

The unused AVAILABLE_SIMD_LEVELS_A2 is deleted. It is the NEON-to-NONE trap in constant form, with no users.

This change keeps the existing mask names. A follow-up renames them to say what they hold.

Out of scope

In rough order of remaining value:

• IndexPQ.cpp and IndexIVFPQ.cpp still pin PQ to ARM_NEON on an SVE host. A mask change is not enough. pq_code_distance-sve.cpp includes only pq_scan_impl.h, and with_HammingComputer<ARM_SVE> has no complete type.
• distances_aarch64.cpp forwards five ARM_NEON specializations to <SIMDLevel::NONE>, so a non-SVE aarch64 host runs scalar code in the IVFFlat scan.
• rabitq_neon.cpp forwards all six ARM_NEON specializations to <SIMDLevel::NONE>. There is no SVE variant.
• There is no block_l2<ARM_NEON> and no exhaustive_L2sqr_blas_cmax<ARM_NEON>.
• Part two of T287037898, which is SDOT and SMMLA. FEAT_DotProd and FEAT_I8MM are not expressible as a SIMDLevel. The right shape is a runtime getauxval(AT_HWCAP) check inside the ARM translation unit, which follows the SIMDConfig::avx512_split precedent. Note the reporter's caveat: SMMLA shows no improvement until the scan loop is tiled, so the tiling must land in the same change.

Differential Revision: D118682598

[meta-codesync]

@mnorris11 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D118682598.