perf(linalg): let AVX-512 hosts fall back to the AVX2 dist_table kernel - #8866
perf(linalg): let AVX-512 hosts fall back to the AVX2 dist_table kernel#8866LuciferYang wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
❌ 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") => |
There was a problem hiding this comment.
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.
a4de31f to
8e36ba8
Compare
There was a problem hiding this comment.
❌ 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.
What this changes
sum_4bit_dist_table_uninit's AVX2 arm matched onlySimdSupport::Avx2. It now matchesAvx512 | Avx512FP16 | Avx2behind anis_x86_feature_detected!("avx2")guard, which is what its siblingsum_4bit_hacc_dist_table_uninitin the same file already does.Why
SIMD_SUPPORTis a single exclusive tier. The initializer inlance-core/src/utils/cpu.rsisif x86::has_avx512() { Avx512FP16 or Avx512 } else if avx2 && fma { Avx2 }, so an AVX-512 host reportsAvx512orAvx512FP16and neverAvx2. 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 exactlyis_x86_feature_detected!("avx512f"), andavx512fdoes not implyavx512bw. Knights Landing has avx512f/cd/er/pf and AVX2, but nobw.kernel_support = "avx512_dist_table"is only set whenbuild.rscompilesdist_table.cwith-march=sapphirerapids. A compiler without sapphirerapids support leaves it unset with acargo:warningand keeps going.The AVX2 kernel needs nothing beyond what the guard checks:
sum_dist_table_32bytes_batch_avx2is#[target_feature(enable = "avx2")]and its intrinsics are AVX2 integer ops plus_mm256_permute2f128_si256and_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_addand 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 au16accumulator.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_SUPPORTstatic 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 theNeonarm, and thex86_64-apple-darwintarget 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
Avx2tier and the siblinghaccfunction.cargo fmt --all -- --check: cleancargo 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: cleancargo test --profile ci -p lance-linalg --lib simd::dist_table, both targets: 9 passedIf 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.rsanddot.rs. There the AVX-512 arm is cfg'd onkernel_support = "avx512_f16"or"avx512_bf16"and the next arm listsAvx2 | Avx512, omittingAvx512FP16. Withfp16kernelson and a compiler that could not build the AVX-512 C, anAvx512FP16host 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=haswellwhile theAvx512*tiers are not FMA-checked. I will send it separately.