Inline PQ code distance kernels into scanner TUs (#5159)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5159

After the SIMD dispatch refactoring, the PQ code distance implementations
(pq_code_distance_single_impl<SL>, pq_code_distance_four_impl<SL>) lived in
separate translation units from the scanner loops (IVFPQScannerT::scan_list_with_table,
PQDistanceComputer::distance_to_code). The compiler could not inline the SIMD
gather/accumulate code into the hot inner loops.

This diff converts the per-SIMD pq_code_distance .cpp files to .h headers and
includes them in the corresponding scanner TUs before the scanner _impl.h
includes. This puts the kernel definitions in the same TU, enabling the compiler
to inline the AVX2/AVX512 vgatherdps code directly into scan_list_with_table and
scan_list_polysemous_hc.

Changes:
- pq_code_distance-avx2.cpp → pq_code_distance-avx2.h (header, #pragma once)
- pq_code_distance-avx512.cpp → pq_code_distance-avx512.h (same)
- New pq_code_distance-generic.h with inline NONE/ARM_NEON specializations
- Scanner TUs (avx2.cpp, avx512.cpp, neon.cpp) include the PQ distance headers
- PQCodeDistance gains static constexpr simd_level member
- scan_list_polysemous uses PQCodeDist::simd_level for Hamming computer dispatch
  (was hardcoded to SIMDLevel::NONE)
- Scanner TUs include per-ISA hamming_computer headers for SIMD Hamming dispatch
- Build files (xplat.bzl, CMakeLists.txt) updated

Verified via objdump that scan_list_with_table and scan_list_polysemous_hc contain
zero calls to pq_code_distance_*_impl — all AVX2/AVX512 gather code is fully inlined.

Reviewed By: algoriddle

Differential Revision: D102942787

fbshipit-source-id: d2979f9cd629652ac1de4886afd3c335fbb4fac8
This commit is contained in:
Matthijs Douze
2026-05-04 07:35:56 -07:00
committed by meta-codesync[bot]
parent 417c53e0ce
commit 28b2b66cf2
19 changed files with 1953 additions and 1036 deletions
+8 -5
View File
@@ -19,7 +19,11 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/ProductQuantizer.h>
#include <faiss/utils/pq_code_distance.h>
#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
// Pull in inline NONE specializations so the test can call
// PQCodeDistance<PQDecoder8, SIMDLevel::NONE> directly.
#include <faiss/impl/pq_code_distance/pq_code_distance-generic.h>
size_t nMismatches(
const std::vector<float>& ref,
@@ -152,8 +156,8 @@ void test(
for (size_t k = 0; k < 1000; k++) {
#pragma omp parallel for schedule(guided)
for (int64_t i = 0; i < (int64_t)n; i++) {
resultsDispatched1x[i] = faiss::pq_code_distance_single(
subq, 8, lookup.data(), codes.data() + subq * i);
resultsDispatched1x[i] = faiss::pq_code_distance_8bit_single(
subq, lookup.data(), codes.data() + subq * i);
}
}
const auto endingTimepoint = std::chrono::steady_clock::now();
@@ -171,9 +175,8 @@ void test(
for (size_t k = 0; k < 1000; k++) {
#pragma omp parallel for schedule(guided)
for (int64_t i = 0; i < (int64_t)n; i += 4) {
faiss::pq_code_distance_four(
faiss::pq_code_distance_8bit_four(
subq,
8,
lookup.data(),
codes.data() + subq * (i + 0),
codes.data() + subq * (i + 1),