Skip to content

fix(linalg): stop FP16 hosts falling to scalar when the AVX-512 f16 kernel is absent #8868

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LuciferYang wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(linalg): stop FP16 hosts falling to scalar when the AVX-512 f16 kernel is absent #8868

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions rust/lance-linalg/src/distance/cosine.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,16 @@ impl Cosine for bf16 {
SimdSupport::Avx512FP16 => unsafe {
bf16_kernel::cosine_bf16_avx512(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_bf16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {
bf16_kernel::cosine_bf16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
bf16_kernel::cosine_bf16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
Expand Down Expand Up @@ -177,6 +187,16 @@ impl Cosine for f16 {
SimdSupport::Avx512FP16 => unsafe {
kernel::cosine_f16_avx512(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_f16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {
kernel::cosine_f16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)

@lance-gatekeeper lance-gatekeeper Bot Aug 29, 2026

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.

This new branch can read past x from safe Rust. Cosine::cosine_fast does not validate equal lengths, but cosine_f16_avx2 receives y.len() and its C loop reads both x[i] and y[i]. In the base cfg-unset path, cosine_scalar instead reaches the checked Dot::dot and panics. Please validate the lengths before dispatch in both the f16 and bf16 implementations so every unsafe C arm inherits the boundary check.

Reproducer run on this head

I forced both AVX-512 builds to fail with a CC wrapper that rejects -march=sapphirerapids, then added a test which mmaps two pages, protects the second with PROT_NONE, places a one-element f16 x at the end of the first page, and calls <f16 as Cosine>::cosine_fast(x, 1.0, &[1.0, 2.0]).

CC=/path/to/cc-no-spr CARGO_TARGET_DIR=/home/agent/tmp/gate-8868-target cargo test --profile ci -p lance-linalg --test gate_8868_cosine_oob --features fp16kernels -- --nocapture
... process didn't exit successfully ... (signal: 11, SIGSEGV: invalid memory reference)

},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
kernel::cosine_f16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
Expand Down
20 changes: 20 additions & 0 deletions rust/lance-linalg/src/distance/dot.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,16 @@ impl Dot for bf16 {
SimdSupport::Avx512FP16 => unsafe {
bf16_kernel::dot_bf16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_bf16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {

@lance-gatekeeper lance-gatekeeper Bot Aug 29, 2026

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.

Checking only fma does not establish the instruction set required by these unsafe calls. Avx512FP16 detection proves only OS-usable AVX-512F plus the FP16 CPUID bit, while these fallback objects are compiled with -march=haswell: the bf16 object emits AVX2 (vpmovzxwd ymm) and the f16 object emits F16C (vcvtph2ps), in addition to FMA. A virtualized or future feature set can therefore select this arm without supporting the called instructions. Please gate bf16 on AVX2+FMA and f16 on F16C+FMA (and audit the complete emitted requirements), or compile against narrower explicit flags and centralize the tested capability predicate.

Verification run on this head 
cc -std=c17 -O3 -ffast-math -funroll-loops -DSUFFIX=_gate -march=haswell -S -o - rust/lance-linalg/src/simd/f16.c | rg 'vcvtph2ps|vfmadd'
cc -std=c17 -O3 -ffast-math -funroll-loops -DSUFFIX=_gate -march=haswell -S -o - rust/lance-linalg/src/simd/bf16.c | rg 'vpmovzxwd|vfmadd'

Both commands produced the named instructions; lance-core/src/utils/cpu.rs does not test F16C or AVX2 when constructing Avx512FP16.

bf16_kernel::dot_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
bf16_kernel::dot_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
Expand Down Expand Up @@ -210,6 +220,16 @@ impl Dot for f16 {
SimdSupport::Avx512FP16 => unsafe {
kernel::dot_f16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_f16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {
kernel::dot_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
kernel::dot_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
Expand Down
20 changes: 20 additions & 0 deletions rust/lance-linalg/src/distance/l2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,16 @@ impl L2 for bf16 {
SimdSupport::Avx512FP16 => unsafe {
bf16_kernel::l2_bf16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_bf16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {
bf16_kernel::l2_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
bf16_kernel::l2_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
Expand Down Expand Up @@ -242,6 +252,16 @@ impl L2 for f16 {
SimdSupport::Avx512FP16 => unsafe {
kernel::l2_f16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_f16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {
kernel::l2_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
kernel::l2_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32)
Expand Down
20 changes: 20 additions & 0 deletions rust/lance-linalg/src/distance/norm_l2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,16 @@ impl Normalize for f16 {
SimdSupport::Avx512FP16 => unsafe {
kernel::norm_l2_f16_avx512(vector.as_ptr(), vector.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_f16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {
kernel::norm_l2_f16_avx2(vector.as_ptr(), vector.len() as u32)

@lance-gatekeeper lance-gatekeeper Bot Aug 29, 2026

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 newly selected f16 AVX2 path does not satisfy the repository's existing numerical tests on the exact target configuration. On this Sapphire Rapids host, forcing the AVX-512 builds to fail made the full library suite execute these cfg-unset arms: test_l2_distance_f16 and test_l2_norm_f16 failed their max_relative = 1e-6 assertions (450 passed; 2 failed; 1 ignored). Replaying the persisted cases with fp16kernels disabled took the prior scalar route and both passed. Please make the fallback meet the existing numerical contract, or otherwise establish and test an explicitly accepted result contract before routing these hosts to it.

Reproducer

The wrapper delegates to cc except that it exits unsuccessfully for -march=sapphirerapids, leaving kernel_support="avx512_f16" and kernel_support="avx512_bf16" genuinely unset.

CC=/path/to/cc-no-spr CARGO_TARGET_DIR=/home/agent/tmp/gate-8868-target cargo test --profile ci -p lance-linalg --lib --features fp16kernels

Observed examples were L2 439198480000.0 versus 439199070000.0, and norm-L2 621139.2 versus 621139.8.

},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
kernel::norm_l2_f16_avx2(vector.as_ptr(), vector.len() as u32)
Expand Down Expand Up @@ -115,6 +125,16 @@ impl Normalize for bf16 {
SimdSupport::Avx512FP16 => unsafe {
bf16_kernel::norm_l2_bf16_avx512(vector.as_ptr(), vector.len() as u32)
},
#[cfg(all(
feature = "fp16kernels",
not(kernel_support = "avx512_bf16"),
target_arch = "x86_64"
))]
// `Avx512*` tier detection does not check FMA, and these kernels are
// compiled `-march=haswell`, which enables it.
SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe {
bf16_kernel::norm_l2_bf16_avx2(vector.as_ptr(), vector.len() as u32)
},
#[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
bf16_kernel::norm_l2_bf16_avx2(vector.as_ptr(), vector.len() as u32)
Expand Down
Loading