faiss DD: use 256-bit fast-scan QBS kernel on AMD Zen 4 (split AVX-512) (#5488)

Summary:
Pull Request resolved: #5488

## TLDR make PQFS faster on AMD Bergamo specifically, without regressing for RaBitQFS, by changing AVX512 --> AVX2 for the FastScan part. The RaBitQ avx512 popcount is still avx512, it is just the FastScan part that stays on AVX2.

Under dynamic dispatch (`faiss.dynamic_dispatch=true`) the fast-scan QBS search
path picks the 512-bit accumulate kernel whenever the runtime SIMD level is
AVX-512. On AMD Zen 4 / Zen 4c ("Bergamo", family 0x19) that is the wrong
choice: Zen 4 splits 512-bit ops over two separate 256-bit parts.
- This can execute operations in 2 separate operations, so 2 cycles compared to a single cycle in avx2. It can sometimes be slower or faster depending on the operation, because more new instructions are supported in avx512 like vpopcnt which are completely missing in AVX2.
- Why split at all? because older CPUs had to downclock for avx512 instructions, while this approach does not, so it can be faster for certain workloads, and it allows support for avx512 instructions for free.

Captured from benchmark infra, matched
AVX2-static vs AVX-512-DD operating points):

- Bergamo `PQ8x4fs`: median AVX-512/AVX2 = 0.861 (nq=1) / 0.860 (batched)
  across 15 datasets, spread [0.857, 0.867], 14/15 datasets below 0.95
  (worst `sift-1M` = 0.60).
- Bergamo `PQ16x4fs`: 0.870 (nq=1) / 0.878 (batched) across 14 datasets.

The magnitude is tight and reproduces on ~every dataset in both the single-query
and batched regimes, well below the Bergamo benchmark noise floor (SVS-Vamana
negative control median ~1.05), so it is real signal. It is isolated to the
bbs=32 QBS path: `PQ*x4fs_64` (bbs=64) and `IVF*,PQ*x4fs` are unaffected, and on
Intel Skylake / Cooper Lake the same factories are FASTER under AVX-512
(1.01-1.71x), so the fix must be keyed to Zen 4 only.

Fix (3 files):
- `simd_levels.{h,cpp}`: add `SIMDConfig::avx512_split`, set by raw-CPUID
  detection (vendor == AuthenticAMD && display_family == 0x19), run once at load
  time from `auto_detect_simd_level()` in both DD and static builds.
- `dispatching.h` `ScannerMixIn::accumulate_loop_qbs`: this is the live QBS
  search dispatch for BOTH PQ fast-scan and RaBitQ fast-scan. When the CPU splits avx512, route to the 256-bit (AVX2) QBS kernel
  (`pq4_accumulate_loop_qbs_fixed_scaler_256<AVX2>`, the same one the existing
  unknown-qbs fallback already uses) instead of the 512-bit kernel. Process-
  constant runtime branch, hoisted out of the inner accumulate loop. Intel
  AVX-512 keeps the 512-bit kernel unchanged.

Note: the fix is placed in `dispatching.h` (the path IndexFastScan /
IndexIVFFastScan search actually take, via `ScannerMixIn`), NOT in
`decompose_qbs.h` -- the QBS block kernel there is only reached by the
`accumulate_to_mem` test utility, not by any search path, so routing it there
would be inert for search.

Scope: RaBitQ fast-scan shares this same PQ4 QBS accumulate, so it is also routed
to 256-bit on Zen 4. That is not expected to regress it -- the 256-bit kernel is
>= the 512-bit one on Zen 4 (the whole premise), and RaBitQ's AVX-512 popcount
(folded into the query LUT) and multibit FP-refine kernels are untouched. An
earlier version of this diff also added a Cooper Lake RaBitQ-popcount cap and a
native `vpopcntq` path; the Cooper Lake regression that motivated them did not
survive additional benchmark runs, so both were
dropped to keep this a single-purpose, evidence-backed fix.

Validated on Bergamo (see Test Plan)

Reviewed By: junjieqi

Differential Revision: D113617130

fbshipit-source-id: c59b522134f40cc7a7daf6abaa3c3f6d5b139fe1

Author: Michael Norris

faiss/impl/fast_scan/dispatching.h

@@ -33,6 +33,7 @@
 
 #include <faiss/impl/fast_scan/accumulate_loops.h>
 #include <faiss/impl/fast_scan/fast_scan.h>
+#include <faiss/utils/simd_levels.h>
 
 #if defined(COMPILE_SIMD_AVX512) && defined(__AVX512F__)
 #include <faiss/impl/fast_scan/accumulate_loops_512.h>
@@ -114,8 +115,40 @@ struct ScannerMixIn : FastScanCodeScanner {
         constexpr bool use_avx512_qbs = false;
 #endif
         if constexpr (use_avx512_qbs) {
-            // Use 512-bit QBS kernels with properly-leveled scalers.
-            if (pq2x4_scale) {
+            // AMD Zen 4 / Zen 4c ("Bergamo", family 0x19) split 512-bit
+            // ops over a 256-bit datapath, so the 512-bit QBS kernel yields no
+            // throughput gain but pays extra per-block LUT-assembly and
+            // cross-lane reduction overhead (measured ~14% search regression
+            // for PQ8x4fs / PQ16x4fs). Route those CPUs to the 256-bit (AVX2)
+            // QBS kernel instead -- same output, no downside on Zen 4. This is
+            // a process-constant runtime branch, hoisted out of the inner
+            // accumulate loop. Intel AVX-512 keeps the 512-bit kernel.
+            if (SIMDConfig::avx512_split) {
+                if (pq2x4_scale) {
+                    NormTableScaler<SIMDLevel::AVX2> scaler(pq2x4_scale);
+                    pq4_accumulate_loop_qbs_fixed_scaler_256<SIMDLevel::AVX2>(
+                            qbs,
+                            nb,
+                            nsq,
+                            codes,
+                            LUT,
+                            handler_,
+                            scaler,
+                            block_stride);
+                } else {
+                    DummyScaler<SIMDLevel::AVX2> dummy;
+                    pq4_accumulate_loop_qbs_fixed_scaler_256<SIMDLevel::AVX2>(
+                            qbs,
+                            nb,
+                            nsq,
+                            codes,
+                            LUT,
+                            handler_,
+                            dummy,
+                            block_stride);
+                }
+            } else if (pq2x4_scale) {
+                // Use 512-bit QBS kernels with properly-leveled scalers.
                 NormTableScaler<THE_LEVEL_TO_DISPATCH> scaler(pq2x4_scale);
                 pq4_accumulate_loop_qbs_fixed_scaler_512(
                         qbs,

faiss/utils/simd_levels.cpp

@@ -20,6 +20,10 @@ SIMDLevel SIMDConfig::level = SIMDLevel::NONE;
 // Bitmask of supported SIMD levels (1 << SIMDLevel)
 uint64_t SIMDConfig::supported_simd_levels = 0;
 
+// Microarchitecture flags (x86). Default false; set by
+// detect_x86_uarch_flags() at load time.
+bool SIMDConfig::avx512_split = false;
+
 // ARM SVE runtime detection
 #if defined(__aarch64__) || defined(_M_ARM64)
 
@@ -53,6 +57,43 @@ static bool has_sve() {
 }
 #endif
 
+// Detect x86 microarchitecture flags used for kernel routing. Uses raw
+// cpuid so it is safe to run on any CPU regardless of compiled SIMD level.
+#if defined(__x86_64__)
+namespace {
+void detect_x86_uarch_flags() {
+    unsigned int eax, ebx, ecx, edx;
+
+    // Vendor string (CPUID.0): "AuthenticAMD" is EBX="Auth", EDX="enti",
+    // ECX="cAMD".
+    eax = 0;
+    ecx = 0;
+    asm volatile("cpuid"
+                 : "=a"(eax), "=b"(ebx), "=c"(ecx), "=d"(edx)
+                 : "a"(eax), "c"(ecx));
+    const bool is_amd =
+            ebx == 0x68747541u && edx == 0x69746e65u && ecx == 0x444d4163u;
+
+    // Family/model (CPUID.1 EAX).
+    eax = 1;
+    ecx = 0;
+    asm volatile("cpuid"
+                 : "=a"(eax), "=b"(ebx), "=c"(ecx), "=d"(edx)
+                 : "a"(eax), "c"(ecx));
+    const unsigned int base_family = (eax >> 8) & 0xfu;
+    const unsigned int display_family =
+            base_family + (base_family == 0xfu ? ((eax >> 20) & 0xffu) : 0u);
+    // AMD Zen 4 / Zen 4c (Bergamo) is family 0x19 and splits AVX-512.
+    // (Zen 5, family 0x1A, has a native 512-bit datapath and is excluded.)
+    SIMDConfig::avx512_split = is_amd && display_family == 0x19u;
+}
+} // namespace
+#else
+namespace {
+void detect_x86_uarch_flags() {}
+} // namespace
+#endif
+
 #ifdef FAISS_ENABLE_DD
 
 // =============================================================================
@@ -101,6 +142,8 @@ bool SIMDConfig::is_simd_level_available(SIMDLevel l) {
 SIMDLevel SIMDConfig::auto_detect_simd_level() {
     SIMDLevel detected_level = SIMDLevel::NONE;
 
+    detect_x86_uarch_flags();
+
 #if defined(__x86_64__) && \
         (defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_AVX512))
     unsigned int eax, ebx, ecx, edx;
@@ -264,6 +307,7 @@ bool SIMDConfig::is_simd_level_available(SIMDLevel l) {
 }
 
 SIMDLevel SIMDConfig::auto_detect_simd_level() {
+    detect_x86_uarch_flags();
     // In static mode, return the compiled-in level
 #if defined(COMPILE_SIMD_AVX512_SPR)
     return SIMDLevel::AVX512_SPR;

faiss/utils/simd_levels.h

@@ -161,6 +161,12 @@ struct FAISS_API SIMDConfig {
     /// Returns bitmask of supported SIMD levels (1 << SIMDLevel).
     static uint64_t supported_simd_levels;
 
+    /// CPU implements AVX-512 by splitting over a 256-bit datapath
+    /// (AMD Zen 4 / Zen 4c "Bergamo", family 0x19). On such CPUs 512-bit
+    /// ops give no throughput gain, so the fast-scan QBS path prefers the
+    /// 256-bit kernel.
+    static bool avx512_split;
+
     static SIMDLevel auto_detect_simd_level();
 
     static constexpr bool has_dynamic_dispatch() {