Skip to content

Optimize quantized and float32 distance kernels - #99

Merged
gouyt13 merged 3 commits into
mainfrom
simd-kernel-optimizations
Sep 12, 2026
Merged

Optimize quantized and float32 distance kernels#99
gouyt13 merged 3 commits into
mainfrom
simd-kernel-optimizations

Conversation

@gouyt13

@gouyt13 gouyt13 commented Sep 12, 2026

Copy link
Copy Markdown
Member

Problem and change

Several distance kernels serialize accumulation through one dependency chain, and raw float32 distances rely on Eigen instead of the library's runtime SIMD dispatch. This change introduces independent accumulators and shared handwritten float32 kernels while preserving distance conventions and packed-code layouts.

  • Update AVX2 extra-code inner products across 1–8 bits, using independent accumulators and balanced reductions. Update AVX-512 1/2/4/6/8-bit kernels; retain the existing 3/5/7-bit implementations.
  • Replace AVX2 binary-dot bit reversal with per-lane shifts and masked loads, in both the standalone kernel and HNSW's inline implementation.
  • Route float32 squared L2, dot, inner-product distance (1 - dot), and squared norm through generic/AVX2/AVX-512 dispatch. Share four-accumulator SIMD implementations, with masked tails for unpadded and unaligned inputs. Generic float32 kernels use scalar loops with double accumulation; non-float templates retain Eigen.
  • Add reference, alignment, bit-order, cancellation, and guard-page coverage. Extend Python IVF raw-reranking/persistence tests to dimensions 65, 420, and 960 for both metrics.

Validation

Checked the merged branch at 2377a92 with native optimization disabled:

  • cmake --build /tmp/rapdx-kernel-tests -j 4: passed.
  • ctest --test-dir /tmp/rapdx-kernel-tests --output-on-failure -j 4: 78 passed, including explicit AVX2 and AVX-512 backend coverage.
  • Rebuilt the Python extension, verified the imported artifact's SHA-256 against the build output, then ran the existing environment's python -m pytest tests/python -q: 118 passed.
  • Repository formatting, Python lint/format, and diff whitespace checks passed.
  • Full clang-tidy passed for all configured first-party translation units. Analysis used temporary include-path adjustments for the installed LLVM OpenMP header and the build environment's pybind11 headers; installed headers and repository tooling were unchanged.

Earlier local production-kernel benchmarks used a Xeon Gold 6418H, one thread pinned to CPU 0, GCC 14.3 (also Clang 15 for binary-dot/float studies), -O3 -DNDEBUG, explicit ISA flags, and repeated alternating baseline/candidate timings. Ratios below are baseline time / candidate time for cached, aligned inputs:

  • AVX2 packed-code kernels at dimension 1024: 1.23–1.78× for 2–8 bits; 1 bit was approximately neutral (1.03×). Fifteen timing pairs, 20,000 warmup and 200,000 timed calls per sample, one query and 64 packed vectors, seed 20260911.
  • AVX-512 updated packed-code kernels at dimension 4096: 1.14–2.21× across 1/2/4/6/8 bits. Two sessions of nine alternating pairs; hot/ring64 workloads and aligned/unaligned layouts. At dimension 128, some adopted kernels regress by about 5–6%.
  • AVX2 binary dot at dimension 1024: 2.56× GCC / 2.63× Clang. Two shuffled sessions, 18 paired observations per comparison, with hot/ring64/32-MiB traversal workloads.
  • Clang AVX2 raw L2 versus same-ISA Eigen: 1.14× at dimension 420, 0.96× at 960 for cached aligned inputs. The latter regression is accepted to use the same handwritten implementation across compilers.

These are prior local microbenchmarks, not measurements rerun after the merge or a universal whole-index speedup claim. Host load/frequency was uncontrolled; performance depends on dimensions, cache behavior, compiler, and CPU. Benchmark artifacts remain outside this branch in the local kernel-study workspace.

Compatibility

No existing public API signatures, index formats, packed-code layouts, ISA requirements, or production dependencies change. Runtime dispatch retains a generic float32 fallback. Floating-point reassociation/FMA can change rounding and near-tie ordering; reference tests use numerical tolerances rather than bitwise equality.

- Shorten accumulation chains in AVX2 and AVX-512 extra-code kernels.
- Optimize AVX2 binary dot products, including the HNSW search path.
- Add shared float32 L2, dot, IP, and norm kernels with runtime SIMD
  dispatch, masked tails, and scalar fallbacks without Eigen.
- Preserve distance conventions, packed layouts, and non-float templates.
- Add kernel regression tests and IVF coverage at dimensions 420 and 960.
@gouyt13 gouyt13 self-assigned this Sep 12, 2026
@gouyt13
gouyt13 merged commit 95aa1de into main Sep 12, 2026
17 checks passed
@gouyt13
gouyt13 deleted the simd-kernel-optimizations branch September 12, 2026 13:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant