Skip to content

[Feature]: Add AVX2/AVX512 SIMD Kernels for Distance in Turbo #706

Description

@richyreachy

Problem / Motivation

Currently, Turbo does not support AVX-related acceleration kernels. To make the library fully featured, we need to implement AVX2 and AVX-512 optimized kernels for distance computations (including SquaredEuclidean, InnerProduct, and Cosine). These new kernels should also be registered in the kernel dispatch table.

Proposed Solution

  1. Kernel source tree reorganized by architecture × quantizer family
    New layout under src/turbo/distance/:
    avx2/, avx512/, avx512_vnni/, neon/, scalar/ × fp32/, fp16/, record_quantized_int8/, record_quantized_int4/, raw_uint8/, pq_quantizer_int8/
    One metric per file (inner_product.cc, squared_euclidean.cc, cosine.cc); shared per-family helpers in common.h
    Unified fp16 availability macros in distance/common/fp16_common.h
  2. Unified kernel dispatch (turbo.cc / turbo.h)
    All kernels registered in a single constexpr KernelSet kKernelTable[]; row order encodes priority (SIMD rows before scalar fallbacks)
    Runtime CPUID gating (CpuSupports + fine-grained feature masks: AVX512BW/DQ, F16C)
    New aggregate lookup get_distance_kernels() returns {dist, batch, preprocess} as one kernel family, preventing cross-family mispairing (e.g. VNNI batch kernels require the +128 uint8-shifted query produced by their paired preprocess)
    Returns empty DistanceKernels{} when no kernel matches the requested architecture
  3. One-to-many batch kernels
    Every batch distance function is now a genuine 1×N kernel (templated batch_impl<dp_batch>, batch_size = 2, prefetch_step = 2): query blocks loaded/converted once per iteration and reused across candidates, with software prefetch of upcoming candidates
    Record-quantized int8/int4 batches share raw_inner_product_batch templates; metric-specific tail math (scale/bias/sum/sum²) applied per record in closed form — codes are never dequantized back to float
  4. New / updated kernels
    AVX512-VNNI: raw uint8 + fp16 squared euclidean; int8 kernels with +128 query preprocess (shift_int8_to_uint8_avx512)
    AVX2/AVX512 fp32 & fp16: inner product / squared euclidean / cosine (cosine delegates to IP)
    NEON: FHT rotate and PQ int8 distance updates
    Unsupported-ISA #else branches standardized to (void) casts (never silent scalar fallback)
  5. Tests
    New per-quantizer test files (tests/turbo/turbo_{fp32,fp16,int8,int4}_quantizer_test.cc) with per-CPU-architecture SIMD distance validation in separate functions
    Batch kernels validated against single-pair kernels across many dims (incl. unaligned) and batch sizes

Alternatives Considered

No response

Affected Area

No response

Metadata

Metadata

Assignees

Labels

featureNew feature wanted

Type

No type

Projects

Status
Backlog

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions