Skip to content

perf(linalg): let AVX-512 hosts fall back to the AVX2 dist_table kernel - #8866

Open
LuciferYang wants to merge 1 commit into
lance-format:mainfrom
LuciferYang:perf/dist-table-avx2-fallback
Open

perf(linalg): let AVX-512 hosts fall back to the AVX2 dist_table kernel#8866
LuciferYang wants to merge 1 commit into
lance-format:mainfrom
LuciferYang:perf/dist-table-avx2-fallback

Conversation

@LuciferYang

Copy link
Copy Markdown
Contributor

What this changes

sum_4bit_dist_table_uninit's AVX2 arm matched only SimdSupport::Avx2. It now matches Avx512 | Avx512FP16 | Avx2 behind an is_x86_feature_detected!("avx2") guard, which is what its sibling sum_4bit_hacc_dist_table_uninit in the same file already does.

Why

SIMD_SUPPORT is a single exclusive tier. The initializer in lance-core/src/utils/cpu.rs is if x86::has_avx512() { Avx512FP16 or Avx512 } else if avx2 && fma { Avx2 }, so an AVX-512 host reports Avx512 or Avx512FP16 and never Avx2. That made the AVX2 arm unreachable on AVX-512 hardware, and since the _ arm is scalar, any AVX-512 host that missed the arm above ran the scalar kernel while holding perfectly good AVX2.

Two ways to miss the arm above, both real:

  • has_avx512() is exactly is_x86_feature_detected!("avx512f"), and avx512f does not imply avx512bw. Knights Landing has avx512f/cd/er/pf and AVX2, but no bw.
  • kernel_support = "avx512_dist_table" is only set when build.rs compiles dist_table.c with -march=sapphirerapids. A compiler without sapphirerapids support leaves it unset with a cargo:warning and keeps going.

The AVX2 kernel needs nothing beyond what the guard checks: sum_dist_table_32bytes_batch_avx2 is #[target_feature(enable = "avx2")] and its intrinsics are AVX2 integer ops plus _mm256_permute2f128_si256 and _mm256_storeu_si256. No FMA, no gather, no VNNI.

One behaviour change worth naming

On the hosts this newly covers, overflow goes from saturating to wrapping. The scalar kernel accumulates with saturating_add and the AVX2 kernel wraps, a divergence the file already documents where the scalar reference is defined. This aligns those hosts with every other AVX2 host rather than introducing anything new, but it is a result change on inputs that overflow a u16 accumulator.

Test plan

This is not testable as it stands, and I would rather say so than imply the green runs below are evidence. The match reads the SIMD_SUPPORT static directly, so a test cannot present a tier, and neither skip condition can be produced on a host that does not already have it. Neither local run exercises the changed arm: this machine is aarch64 and takes the Neon arm, and the x86_64-apple-darwin target under Rosetta reports no AVX at all and takes the scalar arm.

What the correctness argument rests on is that this is the same kernel, with the same guard, that already serves the Avx2 tier and the sibling hacc function.

  • cargo fmt --all -- --check: clean
  • cargo clippy -p lance-linalg --all-targets -- -D warnings, and the same with --target x86_64-apple-darwin, which is the configuration where both x86 arms compile: clean
  • cargo test --profile ci -p lance-linalg --lib simd::dist_table, both targets: 9 passed

If it is worth making testable, the shape would be to extract the decision into a pure function of the tier and the detected features and parametrize over it, leaving the match to call it.

Follow-up found while reviewing this

The same shape is unfixed at 8 sites in the f16 and bf16 dispatches, in cosine.rs, norm_l2.rs, l2.rs and dot.rs. There the AVX-512 arm is cfg'd on kernel_support = "avx512_f16" or "avx512_bf16" and the next arm lists Avx2 | Avx512, omitting Avx512FP16. With fp16kernels on and a compiler that could not build the AVX-512 C, an Avx512FP16 host matches no arm and runs scalar f16, even though the AVX2 f16 kernel is always built and is a hard build error if it fails. That one needs an FMA guard as well, because those kernels are compiled -march=haswell while the Avx512* tiers are not FMA-checked. I will send it separately.

@github-actions github-actions Bot added A-index Vector index, linalg, tokenizer performance labels Aug 29, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate recommendation: request changes.

The AVX-512-to-AVX2 fallback is correctly ordered and feature-guarded, but this performance fix needs a regression test for the dispatch decision before acceptance. A pure selector exercised with synthetic tiers and feature availability would cover the affected configurations without special hardware.

#[cfg(target_arch = "x86_64")]
SimdSupport::Avx2 => unsafe {
SimdSupport::Avx512 | SimdSupport::Avx512FP16 | SimdSupport::Avx2
if std::arch::is_x86_feature_detected!("avx2") =>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new fallback has no regression test, so removing either Avx512* alternative later would leave every current dist_table test green. On the available AVX-512FP16 + AVX512BW host, cargo test -p lance-linalg --lib simd::dist_table passed all 9 tests by taking the unchanged AVX-512 arm, not this branch. That leaves the exact regression unpinned and conflicts with the repository requirement that every performance fix have a corresponding test. Please extract the fallback eligibility/backend decision into a pure helper and parameterize the AVX-512 and AVX-512FP16 fallback-to-AVX2 cases plus the no-AVX2-to-scalar case.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 29, 2026
@LuciferYang
LuciferYang force-pushed the perf/dist-table-avx2-fallback branch from a4de31f to 8e36ba8 Compare August 29, 2026 16:18
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 29, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate recommendation: request changes.

The rebase leaves the patch and recommendation unchanged. The dispatch-coverage finding remains: the fallback is correctly ordered and feature-guarded, but this performance fix still needs a synthetic tier/feature regression test before acceptance.

A pure selector exercised with AVX-512, AVX-512FP16, and no-AVX2 cases remains the clean path.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-index Vector index, linalg, tokenizer K-changes Latest Gatekeeper recommendation requests changes. performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant