Repository navigation
Add AVX-512 f32 dot product and squared L2 kernels - #145674
Conversation
vec_dotf32 and vec_sqrf32 currently use AVX2 (256-bit ymm registers) in vec_1.cpp. Add vec_dotf32_2, vec_sqrf32_2, and their bulk/bulk_offsets variants to vec_2.cpp using __m512 / _mm512_fmadd_ps with 4 accumulators and masked tail handling. The runtime dispatch in JdkVectorLibrary already falls back from _2 to _1 suffix, so these are picked up automatically on AVX-512 capable hardware.
|
I have inspected the kernels to see if there is more that can be gained; TL;DR: the advantage over AVX2 is purely due to wider loads -- we are able to load more data per cycle. With 32-bit elements, these simple kernels are load-bound, not compute bound. |
|
The current implementation is at (or very near) the theoretical performance limit for a single-pair distance computation on all current x86 microarchitectures. A float32 dot product loads two vectors (a and b) and computes vmovups zmm8, [rdi + rax*4] ; load 16 floats from a
vfmadd231ps zmm3, zmm8, [rsi + rax*4] ; load 16 floats from b, acc += a * bTwo 512-bit loads feed one fused multiply-add. That's 2 loads per FMA. Every x86 CPU — AMD and Intel alike — has at most 2 load ports. A dot product needs 2 loads to produce 1 FMA. This means the load ports are always fully occupied, while the FMA units are at most 50% utilized. No amount of loop unrolling, accumulator count tuning, or instruction scheduling can change this — the ratio is dictated by the algorithm, not the implementation. Evidence from benchmarks (Zen 5, c8a.xlarge)Three independent observations confirm the load-bound diagnosis:
This holds across all current x86 CPUs:
On Zen 5 and other AVX-512 capable processors, load bandwidth doubled alongside register width:
The wider loads process more data per cycle, so we see a speedup wrt AVX2, but the load ports are still the bottleneck in both cases -- this is the most amount of data that can be moved through these CPUs. |
ldematte
left a comment
There was a problem hiding this comment.
One minor thing, and requires a version bump, but otherwise this LGTM
|
Pinging @elastic/es-search-relevance (Team:Search Relevance) |
|
Hi @ChrisHegarty, I've created a changelog YAML for you. |
The
vec_dotf32andvec_sqrf32functions currently use AVX2 (256-bit ymm registers) in vec_1.cpp. This PR addsvec_dotf32_2,vec_sqrf32_2, and their bulk/bulk_offsets variants tovec_2.cppusing __m512 / _mm512_fmadd_ps with 4 accumulators and masked tail handling. The runtime dispatch in JdkVectorLibrary already falls back from_2to_1suffix, so these are picked up automatically on AVX-512 capable hardware.