Native SIMD: Faiss Distance Kernels in C++

The Lucene SIMD chapter told one half of the vectorization story: the lucene engine computes distances in Java, and the Panama Vector API turns a lane-wise FloatVector loop into AVX-512/AVX2/NEON instructions at JIT time. This chapter is the other half — the one the faiss (and legacy nmslib) engine actually runs. Faiss is C++. Its distance kernels are hand-written with SIMD intrinsics (_mm256_*, _mm512_*, ARM vfmaq_f32) and compiled into the per-architecture shared libraries the k-NN plugin loads over JNI. Same hardware instructions as the Panama path; different language, different build, different place to put a breakpoint.

If you ever pick up a faiss-side k-NN issue — a recall regression after a build-flag change, a segfault in a SIMD frame in hs_err_pid, a "why is AVX-512 not being used on this Graviton box" ticket — you need to read C++ SIMD, not just Panama. This chapter is the deliberate C++ mirror of the Java one. It assumes you have read the Panama chapter (so you already know why distance is the hot loop and what a scalar tail is) and the native integration chapter (so you know how the .so gets loaded). Here we go beneath the JNI boundary into the kernels themselves.

After this chapter you can: write a dot-product and an L2 kernel in AVX2, AVX-512, and NEON intrinsics; explain the horizontal-reduce idiom; explain the PQ fast-scan trick where _mm256_shuffle_epi8 becomes a 32-way parallel codebook lookup; reason about whether a distance kernel is compute-bound or memory-bandwidth-bound (and why that tells you when quantization buys latency, not just RAM); and find/verify which SIMD variant faiss actually selected at runtime.

Note: Nothing here is OpenSearch-specific math — it is how faiss (and any serious ANN library) implements the inner loop. OpenSearch inherits it by vendoring faiss under jni/ and compiling it with the right flags. The exact faiss file/function names move between versions; every code location below is a grep/find you run against your own jni/ tree, never a fixed line.

Two SIMD stories, one set of instructions

 k-NN engine = lucene                    k-NN engine = faiss / nmslib
 ───────────────────────                 ─────────────────────────────
 VectorUtil.dotProduct (Java)            faiss fvec_inner_product (C++)
        │                                        │
 Panama FloatVector loop                 _mm256_fmadd_ps intrinsics
        │  (JIT intrinsifies)                    │  (compiler emits)
        ▼                                        ▼
   vfmadd231ps  ymm/zmm   ◄── SAME CPU INSTRUCTION ──►   vfmadd231ps  ymm/zmm

Both engines end at the identical AVX2 vfmadd231ps (or AVX-512 on zmm, or NEON fmla). The difference is who emits it: the HotSpot JIT for the Java path, the C++ compiler (g++ -mavx2) for the native path. This matters operationally because the two are enabled by completely different switches — the Java path needs --add-modules jdk.incubator.vector on the JVM; the native path needs the right faiss library variant to have been built and selected. An operator can have SIMD on for the lucene engine and off (running a generic scalar faiss build) for the faiss engine on the same node. The Panama chapter covered switch #1; this chapter covers switch #2.

# In a k-NN checkout, the native distance kernels live in the vendored faiss + the jni glue:
find jni -path '*faiss*' \( -name 'distances*.cpp' -o -name 'distances_simd*.cpp' \
  -o -name 'ProductQuantizer*.cpp' -o -name 'pq4_fast_scan*.cpp' \)
grep -rn "fvec_inner_product\|fvec_L2sqr\|_mm256_\|_mm512_\|vfmaq_f32" jni 2>/dev/null | head

The dot product in C++ intrinsics

Start from the scalar kernel — the C++ twin of dotScalar from the Panama chapter:

// Scalar: one multiply-add per element. The baseline the SIMD versions must beat
// and must agree with (within float tolerance — FP add is not associative).
float dot_scalar(const float* a, const float* b, size_t n) {
    float sum = 0.f;
    for (size_t i = 0; i < n; ++i) sum += a[i] * b[i];
    return sum;
}

AVX2 — 8 floats per instruction (256-bit ymm)

#include <immintrin.h>   // x86 intrinsics: __m256, _mm256_*

// AVX2 dot product. Processes 8 floats per iteration into a SIMD accumulator,
// then horizontally reduces. Mirrors the Panama dotVectorized exactly.
float dot_avx2(const float* a, const float* b, size_t n) {
    __m256 acc = _mm256_setzero_ps();             // 8 lanes of 0.0f
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 va = _mm256_loadu_ps(a + i);       // unaligned load of 8 floats
        __m256 vb = _mm256_loadu_ps(b + i);
        acc = _mm256_fmadd_ps(va, vb, acc);       // acc += va * vb  (one FMA, one rounding)
    }
    // Horizontal sum of the 8 lanes in `acc` -> one float.
    __m128 lo  = _mm256_castps256_ps128(acc);     // lanes 0..3
    __m128 hi  = _mm256_extractf128_ps(acc, 1);   // lanes 4..7
    __m128 s   = _mm_add_ps(lo, hi);              // 4 partial sums
    s = _mm_add_ps(s, _mm_movehl_ps(s, s));       // fold high 2 into low 2
    s = _mm_add_ss(s, _mm_shuffle_ps(s, s, 0x1)); // fold lane 1 into lane 0
    float sum = _mm_cvtss_f32(s);
    for (; i < n; ++i) sum += a[i] * b[i];        // scalar tail for n % 8
    return sum;
}

The structure is identical to the Panama version, lane for lane:

Note: _mm256_loadu_ps is the unaligned load. The aligned form _mm256_load_ps is a hair faster but faults if a+i is not 32-byte aligned — and vectors arriving from a memory-mapped .faiss file or a JNI buffer are not guaranteed aligned. Faiss uses unaligned loads in the general kernels for exactly this reason. Reaching for _mm256_load_ps on an unaligned pointer is a classic SIGSEGV-in-a-SIMD-frame bug.

AVX-512 — 16 floats per instruction (512-bit zmm)

// AVX-512: double the lane width, plus a dedicated horizontal-reduce intrinsic.
float dot_avx512(const float* a, const float* b, size_t n) {
    __m512 acc = _mm512_setzero_ps();             // 16 lanes
    size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 va = _mm512_loadu_ps(a + i);
        __m512 vb = _mm512_loadu_ps(b + i);
        acc = _mm512_fmadd_ps(va, vb, acc);
    }
    float sum = _mm512_reduce_add_ps(acc);        // 16-lane horizontal sum in one call
    for (; i < n; ++i) sum += a[i] * b[i];        // scalar tail for n % 16
    return sum;
}

AVX-512 gives twice the lanes (16 vs 8) and a tidy _mm512_reduce_add_ps so you do not hand-roll the shuffle cascade. The catch is real-world: AVX-512 can lower the CPU's clock frequency under heavy use (historically on some Intel parts), and it is absent on most ARM and on AVX2-only x86. So a generic build cannot assume it — which is why faiss ships multiple variants and picks at runtime (below).

NEON — 4 floats per instruction (ARM, e.g. AWS Graviton)

OpenSearch runs heavily on ARM (Graviton) hosts, so the native kernels have a NEON path:

#include <arm_neon.h>    // ARM NEON intrinsics: float32x4_t, vfmaq_f32, ...

float dot_neon(const float* a, const float* b, size_t n) {
    float32x4_t acc = vdupq_n_f32(0.f);           // 4 lanes of 0.0f
    size_t i = 0;
    for (; i + 4 <= n; i += 4) {
        float32x4_t va = vld1q_f32(a + i);        // load 4 floats
        float32x4_t vb = vld1q_f32(b + i);
        acc = vfmaq_f32(acc, va, vb);             // acc += va * vb
    }
    float sum = vaddvq_f32(acc);                  // AArch64 horizontal add of 4 lanes
    for (; i < n; ++i) sum += a[i] * b[i];        // scalar tail for n % 4
    return sum;
}

NEON is 128-bit (4 float lanes), so the per-instruction win is smaller than AVX-512's 16, but the shape is the same and the speedup over scalar is still large. This is exactly the point the Panama chapter made about SPECIES_PREFERRED resolving to 4 lanes on NEON — here you see the other engine's version of the same accommodation, written by hand.

L2 squared distance and cosine

The other two space_types are the same skeleton with a different inner op. Squared Euclidean subtracts then squares-accumulates:

// L2 squared distance in AVX2: sum of (a_i - b_i)^2.
float l2sqr_avx2(const float* a, const float* b, size_t n) {
    __m256 acc = _mm256_setzero_ps();
    size_t i = 0;
    for (; i + 8 <= n; i += 8) {
        __m256 d = _mm256_sub_ps(_mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i));
        acc = _mm256_fmadd_ps(d, d, acc);         // acc += d*d
    }
    // ... identical horizontal reduce + scalar tail as dot_avx2 ...
    float s[8]; _mm256_storeu_ps(s, acc);
    float sum = s[0]+s[1]+s[2]+s[3]+s[4]+s[5]+s[6]+s[7];
    for (; i < n; ++i) { float d = a[i]-b[i]; sum += d*d; }
    return sum;
}

Cosine is a dot product plus two norm accumulators (|a|², |b|²) — three FMA chains in the same loop — and a final dot / sqrt(|a|²·|b|²). In practice faiss and OpenSearch prefer to normalize once at index time and use a plain inner product thereafter, because cosine on unit vectors is the dot product. That equivalence is derived in the vector-math foundations chapter; the SIMD consequence is that you almost never run the three-accumulator cosine kernel in the hot path — you run dot_* on pre-normalized data.

Warning: As in Java, do not assert a SIMD kernel equals the scalar kernel exactly. Floating-point addition is not associative, and acc += va*vb accumulates in a different order than the scalar sum += a[i]*b[i]. The results agree within a small relative tolerance; tests assert "within delta," never ==.

The PQ fast-scan trick: _mm256_shuffle_epi8 as a 16-way table lookup

The most beautiful SIMD-meets-algorithm moment in vector search is product-quantization distance, and it is invisible from the Java side. Recall from the algorithms chapter and the quantization math lab how PQ computes a distance with Asymmetric Distance Computation (ADC):

1. Split the d-dim space into m sub-spaces. Each sub-space has a codebook of k = 2^nbits centroids. Each database vector is stored as m centroid ids (one byte each for nbits=8).
2. For a query, precompute a lookup table LUT[m][k]: LUT[s][c] = distance from the query's s-th sub-vector to centroid c of sub-space s.
3. A database vector's approximate distance is then just sum over s of LUT[s][ code[s] ] — m table lookups and m adds, no multiplies, no full vector touched.

Step 3 is a gather — random-indexed loads — which SIMD normally hates. The faiss PQ fast-scan (the pq4 kernels, after André, Kermarrec & Le Scouarnec, "Cache locality is not enough") makes it SIMD-friendly by using nbits = 4 so each codebook has only k = 16 centroids — and 16 bytes is exactly what one _mm256_shuffle_epi8 can index in parallel, per 128-bit lane:

// One sub-space of PQ4 fast-scan, conceptually. `lut` holds 16 uint8 distances
// (one per centroid) broadcast into both 128-bit lanes of a 256-bit register.
// `codes` holds the 4-bit PQ codes of 32 database vectors for THIS sub-space,
// laid out so the shuffle indexes them. One shuffle = 32 codebook lookups.
__m256i lut   = _mm256_loadu_si256((const __m256i*)lut_ptr);     // 16 dists x2 lanes
__m256i codes = _mm256_loadu_si256((const __m256i*)codes_ptr);   // 32 x 4-bit codes
__m256i part  = _mm256_shuffle_epi8(lut, codes);   // 32 parallel table lookups!
acc = _mm256_adds_epu8(acc, part);                 // saturating add into per-vector accumulators

_mm256_shuffle_epi8 treats the low 4 bits of each byte in codes as an index into the 16-byte table in the same 128-bit lane of lut, writing 32 looked-up bytes at once. So one instruction evaluates one sub-space's contribution for 32 database vectors simultaneously. Loop over the m sub-spaces, accumulating with _mm256_adds_epu8 (saturating, to bound the 8-bit range), and you have scored 32 PQ vectors with ~m shuffles instead of 32·m scalar gathers. This is why IVF+PQ / HNSW+PQ with 4-bit fast-scan is dramatically faster per candidate than float32 — not only is each vector 8× smaller (4 bits × m vs 32 bits × d), the distance is a shuffle, not a multiply-reduce.

   one sub-space, 32 database vectors, ONE instruction:

   lut  (16 quantized sub-distances):  [ d0 d1 d2 ... d15 ]   (per 128-bit lane)
   codes (32 4-bit centroid ids):      [ c0 c1 c2 ........ c31 ]
                  _mm256_shuffle_epi8(lut, codes)
   part:                               [ d[c0] d[c1] ... d[c31] ]   <-- 32 lookups, 1 op
   acc += part  (saturating)           accumulate this sub-space into 32 running distances

# Find the fast-scan kernels and the LUT-quantization in the vendored faiss:
grep -rn "shuffle_epi8\|pq4_accumulate\|FastScan\|simd_result_handlers\|lut" \
  $(find jni -path '*faiss*' -name '*fast_scan*' -o -path '*faiss*' -name 'pq4*') 2>/dev/null | head

Note: This is the concrete payoff of a design decision that looks arbitrary on paper — "why would anyone use only 16 centroids per sub-space when 256 gives better recall?" Because 16 fits a shuffle, and the throughput win (often >5×) usually beats the recall loss, which a rescoring pass then recovers. The algorithm shape was chosen for the SIMD instruction. That is the kind of cross-layer reasoning a k-NN maintainer needs.

Compute-bound vs memory-bound: the roofline you must reason about

The Panama chapter said "make the kernel 4× faster and you make search 4× faster." That is true only when the kernel is compute-bound. Often it is memory-bandwidth-bound, and then SIMD FLOPs are not the lever — bytes streamed is. You must know which regime you are in.

Arithmetic intensity of a float32 dot product over d dims:

work    = 2d flops            (d multiplies + d adds)
traffic = 2d * 4 bytes = 8d   (load a[] and b[], 4 bytes each)
intensity = 2d / 8d = 0.25 flops/byte

0.25 flops/byte is very low. A modern core does tens of flops per byte of bandwidth at peak, so a cold dot product — vectors streamed from DRAM, each touched once — is limited by memory bandwidth, not by the FMA unit. AVX-512 cannot help if the bottleneck is getting the bytes to the CPU.

This reframes the whole quantization story:

This is the reason quantization improves latency, not just memory footprint: a PQ vector is 8–32× fewer bytes, so in the memory-bound regime — which is most of large-scale ANN — you stream 8–32× less and the scan gets proportionally faster, on top of the fast-scan shuffle trick. Conversely, in a small cache-resident HNSW graph the win comes from wider lanes and FMA. When you triage "we made the kernel SIMD and it barely got faster," the first question is which regime — and the answer is usually "memory-bound, so go shrink the vectors, not widen the lanes."

# Tell the regime empirically: if perf scales with vector BYTES (float32 -> fp16 -> PQ) more than
# with SIMD width, you're memory-bound. Watch LLC misses and DRAM bandwidth during a scan:
perf stat -e cycles,instructions,LLC-load-misses,fp_arith_inst_retired.256b_packed_single ./bench

How faiss builds and selects the SIMD variant

A single binary cannot just #include <immintrin.h> and call AVX-512 — the instruction faults on a CPU that lacks it. Faiss (and the k-NN jni/ build) handle this with multiple builds + runtime dispatch, the C++ analogue of Lucene's VectorizationProvider:

1. Separate translation units / flags per ISA. The SIMD kernels are compiled in variants — a generic build, an AVX2 build (-mavx2 -mfma), an AVX-512 build (-mavx512f …). On ARM, a NEON build. The k-NN CMake emits correspondingly-suffixed libraries (e.g. a base libopensearchknn_faiss plus AVX2/AVX-512 variants — names are version-specific).
2. Runtime CPU feature detection chooses which to load/call. k-NN has a PlatformUtils-style check (isAVX2SupportedBySystem / isAVX512SupportedBySystem) that the Java side consults when loading the native library, so a node on an AVX2-only CPU loads the AVX2 build and an AVX-512 node loads the AVX-512 build — the same mechanism, different library.

# The CMake that produces the per-ISA libraries, and the dispatch on the Java side:
grep -rn "avx2\|avx512\|simd\|march\|mavx\|add_library\|TARGET_LIB" jni/CMakeLists.txt
grep -rn "isAVX2SupportedBySystem\|isAVX512SupportedBySystem\|PlatformUtils\|loadLibrary\|faiss_avx" \
  src/main/java/org/opensearch/knn 2>/dev/null | head
# Confirm which variants got built:
find . -name 'libopensearchknn_faiss*' -o -name 'libfaiss*'

Warning: -march=native is tempting and wrong for a distributed product. It bakes the build machine's ISA into the binary; ship that to a CPU without those instructions and you get SIGILL (illegal instruction) at the first kernel call — a brutal, late, hard-to-attribute crash. The multi-variant + runtime-dispatch approach exists precisely so the same artifact runs on AVX2, AVX-512, and (separately built) NEON hosts. If you touch the jni/ build, this is the invariant you must not break.

How to verify which kernel is actually running

The native analogue of "is Panama active?" Two questions: does the CPU support the ISA, and did faiss select the matching build?

# 1. What does the CPU support?
grep -o 'avx2\|avx512f\|fma' /proc/cpuinfo | sort -u          # x86
grep -o 'asimd\|neon' /proc/cpuinfo | sort -u                  # ARM

# 2. Which native library got loaded by the JVM (the live JNI boundary)?
PID=$(pgrep -f org.opensearch.bootstrap.OpenSearch | head -1)
grep -E 'libopensearchknn|libfaiss' /proc/$PID/maps | awk '{print $6}' | sort -u
#   look for an avx2/avx512 suffix in the mapped path

# 3. Symbols present in the built library (did the AVX-512 kernels compile in?)
nm -D $(find . -name 'libopensearchknn_faiss*' | head -1) 2>/dev/null \
  | grep -iE 'avx512|avx2|fast_scan' | head

# 4. Microbench / perf-counter the scan and check for 512-bit FP ops actually retiring:
perf stat -e fp_arith_inst_retired.512b_packed_single,fp_arith_inst_retired.256b_packed_single \
  -- <a command that runs a faiss scan>

If /proc/cpuinfo shows avx512f but the mapped library has no AVX-512 suffix and perf shows zero 512b_packed ops, faiss is running a narrower (or scalar) build — the native equivalent of the "silently fell back to the scalar Panama path" bug from the Java chapter.

Common pitfalls

Reading exercise

# In a k-NN checkout with the vendored faiss under jni/:

# 1. The scalar/SIMD distance kernels and their dispatch.
grep -rn "fvec_inner_product\|fvec_L2sqr\|_mm256_fmadd_ps\|_mm512_fmadd_ps\|vfmaq_f32" \
  $(find jni -path '*faiss*' -name 'distances*.cpp') | head

# 2. The PQ fast-scan shuffle.
grep -rn "shuffle_epi8\|pq4\|FastScan\|simd_result_handlers" $(find jni -path '*faiss*') | head

# 3. The per-ISA build + the Java-side dispatch.
grep -rn "avx2\|avx512\|march\|add_library" jni/CMakeLists.txt
grep -rn "isAVX2SupportedBySystem\|isAVX512SupportedBySystem\|loadLibrary" \
  src/main/java/org/opensearch/knn | head

# 4. Which variant loaded on this node.
PID=$(pgrep -f org.opensearch.bootstrap.OpenSearch | head -1)
grep -E 'libopensearchknn|libfaiss' /proc/$PID/maps | awk '{print $6}' | sort -u

Answer:

1. Write, from memory, an AVX2 dot product: the accumulator, the load/FMA loop, the horizontal reduce, the scalar tail. Map each line to its Panama counterpart.
2. Explain how _mm256_shuffle_epi8 evaluates one PQ sub-space for 32 vectors at once, and why nbits=4 (16 centroids) is what makes it possible.
3. Compute the arithmetic intensity of a float32 L2 distance and state whether a cold flat scan is compute- or memory-bound. What does that imply about whether quantization is a latency win?
4. Why does faiss ship multiple per-ISA libraries instead of one -march=native build, and what crash do you get if that invariant is violated?
5. Give two independent ways to verify, on a running OpenSearch node, that the AVX-512 faiss kernels (not a narrower fallback) are actually executing.
6. The lucene engine and the faiss engine both end at vfmadd231ps. Explain the two completely different switches that enable each, and how one can be on while the other is off.

Validation: prove you understand this

1. Implement dot_avx2 and l2sqr_avx2 (or read faiss's) and explain the horizontal-reduce cascade line by line; explain why it is not just acc[0]+acc[1]+... in registers.
2. Explain the PQ fast-scan: the LUT layout, why 4-bit codes, what one shuffle_epi8 produces, and how the saturating accumulate bounds the 8-bit range.
3. Derive arithmetic intensity for dot and L2; state the regime for (a) a cold 100M-vector flat scan and (b) a hot in-cache HNSW walk, and the right optimization for each.
4. Explain faiss's runtime SIMD dispatch and the exact failure mode of -march=native on a mixed fleet.
5. Contrast this chapter's enablement switch with the Panama chapter's: name both, and describe a node where the lucene engine is SIMD-accelerated but the faiss engine is not.
6. Explain why the team chose 16-centroid sub-quantizers for fast-scan even though 256 gives better recall — i.e. how an instruction shaped an algorithm — and what recovers the lost recall.

When you can do all six, you have both halves of the vectorization story — Java/Panama (the Lucene SIMD chapter) and C++ intrinsics (this chapter) — and you can read, fix, and reason about the inner loop on either engine. Next, see how these kernels plug into the compression menu in quantization and disk-based ANN, implement the PQ/SQ/BQ math by hand in Lab VI3, and see the recall cost of all this approximation in Lab VI4: Visualizing ANN.