From 63040e17ccce75f689b3c42071725463cea9851f Mon Sep 17 00:00:00 2001 From: Tobias Perelstein <5562156+tobocop2@users.noreply.github.com> Date: Tue, 28 Apr 2026 00:59:45 -0400 Subject: [PATCH 01/15] feat(lance-linalg): runtime SIMD dispatch for pre-Haswell x86_64 from-source builds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds 5-tier runtime SIMD dispatch (scalar / AVX / AVX+FMA / AVX2+FMA / AVX-512) to the 10 hot f32/f64 distance kernels in lance-linalg, matching the per-tier dispatch shape already used by dot_u8.rs / cosine_u8.rs / l2_u8.rs and the f16/bf16 paths in norm_l2.rs. Today, `import lancedb` SIGILLs on AVX-without-AVX2 CPUs (Intel Sandy Bridge / Ivy Bridge / Westmere; AMD Bulldozer / Piledriver / Steamroller) because the wheel bakes AVX2 into every compiled function with no runtime guard. numpy and pyarrow handle the same hardware via runtime dispatch; this PR brings lance to parity for the from-source legacy build path. Per westonpace's review feedback on lancedb/lancedb#3324 — fast by default, extra steps to work on legacy — the workspace .cargo/config.toml baseline stays at target-cpu=haswell. The qemu-pre-haswell CI job sets RUSTFLAGS=-C target-cpu=x86-64-v2 only in its own env block so the runtime-dispatch path gets exercised under qemu Nehalem without affecting any other build. CONTRIBUTING.md documents the from-source legacy build for users on pre-Haswell hardware. Changes: - lance-core::utils::cpu: extend SimdSupport with Avx and AvxFma tiers, add #[non_exhaustive] for forward compatibility, add SimdInfo + simd_info() introspection API mirroring pyarrow.runtime_info(). - lance-linalg::distance::{cosine, dot, l2, norm_l2}: per-tier dispatch via match *SIMD_SUPPORT for all f32/f64 hot kernels. AVX-512 inner kernels use _mm512_* intrinsics directly; AVX+FMA kernels use the existing f32x8 / f64x8 SIMD types; AVX-only kernels use raw _mm256_mul/_mm256_add intrinsics because f32x8::multiply_add lowers to an FMA instruction. AVX2-host dispatch routes to the AvxFma kernel where the body uses no AVX2-specific intrinsics (Avx2 | AvxFma => x86::FOO_avx_fma) — eliminates ~480 lines of byte-identical AVX2/AvxFma per-tier function pairs across all 10 hot kernels. Documents AvxFma/Avx fall-through to scalar in the u8 and dist_table dispatchers (their AVX2 inners use integer ops not available below AVX2). - lance-linalg::simd::{f32, f64}: drop the dead AVX-512 specializations of f32x16 / f64x8 (gated on target_feature="avx512f" which no project CI configuration enables). Per-tier dispatch in distance/* uses raw _mm512_* intrinsics directly. f32x8::gather adds a runtime AVX2 guard with a scalar fallback for x86_64 hosts without AVX2. - python: expose lance.simd_info() (native binding + re-export in lance/__init__.py with a pytest) so users can verify which SIMD tier the runtime dispatched to without rebuilding. - .github/workflows/rust.yml: new qemu-pre-haswell job that builds with RUSTFLAGS=-C target-cpu=x86-64-v2 and runs lance-linalg lib tests under qemu-x86_64 -cpu Nehalem, catching SIGILL leaks in the legacy-build path before they ship. - CONTRIBUTING.md: documents the legacy build (RUSTFLAGS=-C target-cpu=x86-64-v2 cargo build --release). Zero new external dependencies; Cargo.lock unchanged. lance-linalg lib tests all green (83/83). Verified end-to-end on Sandy Bridge Xeon E5-2609 via the companion lancedb wheel build: pre-PR `pip install lancedb` SIGILLs; post-PR a from-source build with the documented RUSTFLAGS override produces a wheel where import + vector search work at the AVX tier. Closes #6618. --- .github/workflows/rust.yml | 36 + CONTRIBUTING.md | 10 + python/python/lance/__init__.py | 2 + python/python/tests/test_lance.py | 22 + python/src/lib.rs | 24 + rust/lance-core/src/utils/cpu.rs | 178 ++- .../proptest-regressions/distance/cosine.txt | 7 + rust/lance-linalg/src/distance/cosine.rs | 1110 +++++++++++++++-- rust/lance-linalg/src/distance/cosine_u8.rs | 4 + rust/lance-linalg/src/distance/dot.rs | 392 +++++- rust/lance-linalg/src/distance/dot_u8.rs | 4 + rust/lance-linalg/src/distance/l2.rs | 406 +++++- rust/lance-linalg/src/distance/l2_u8.rs | 4 + rust/lance-linalg/src/distance/norm_l2.rs | 341 ++++- rust/lance-linalg/src/simd/dist_table.rs | 4 + rust/lance-linalg/src/simd/f32.rs | 199 +-- rust/lance-linalg/src/simd/f64.rs | 133 +- 17 files changed, 2599 insertions(+), 277 deletions(-) create mode 100644 rust/lance-linalg/proptest-regressions/distance/cosine.txt diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index e3c17671ce3..1bcadeb9199 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -253,6 +253,42 @@ jobs: - name: Check benchmarks run: cargo check --profile ci --benches + qemu-pre-haswell: + # Verifies that lance-linalg's runtime SIMD dispatch still works + # correctly when the binary is built with the lower x86-64-v2 baseline + # (the legacy build path documented in CONTRIBUTING.md). Emulates a + # Nehalem CPU under qemu-user; catches any accidental AVX2/FMA + # instructions that leak past the runtime dispatch. + # + # The published-wheel default baseline (`target-cpu=haswell`) is set in + # `.cargo/config.toml`; this job overrides RUSTFLAGS for one job to + # exercise the legacy path without affecting any other build. + name: pre-Haswell SIGILL check (qemu Nehalem) + runs-on: ubuntu-24.04 + timeout-minutes: 60 + env: + CC: clang + CXX: clang++ + RUSTFLAGS: "-C target-cpu=x86-64-v2" + CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER: "qemu-x86_64 -cpu Nehalem" + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + - name: Setup rust toolchain + run: | + rustup toolchain install stable + rustup default stable + - uses: Swatinem/rust-cache@779680da715d629ac1d338a641029a2f4372abb5 # v2 + - name: Install dependencies + run: | + sudo apt update + sudo apt install -y protobuf-compiler libssl-dev qemu-user + - name: Build lance-linalg lib tests in release mode + run: | + cargo test --release -p lance-linalg --lib --no-run + - name: Run lance-linalg lib tests under qemu Nehalem + run: | + cargo test --release -p lance-linalg --lib + msrv: # Check the minimum supported Rust version name: MSRV Check - Rust v${{ matrix.msrv }} diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8f3ec285f31..d96c7e451fb 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -25,6 +25,16 @@ Currently Lance is implemented in Rust and comes with a Python wrapper. So you'l a. Install pre-commit: https://pre-commit.com/#install b. Run `pre-commit install` in the root of the repo +## Building for legacy x86_64 hosts (pre-Haswell) + +The default workspace build targets `x86-64-haswell` (AVX2 + FMA + F16C), matching the published wheels. To build a binary that runs on pre-Haswell silicon (Sandy Bridge / Ivy Bridge / Westmere on Intel, Bulldozer / Piledriver / Steamroller on AMD — i.e. CPUs without AVX2), set the baseline yourself at build time: + +```sh +RUSTFLAGS="-C target-cpu=x86-64-v2" cargo build --release +``` + +Runtime SIMD dispatch in `lance-linalg::distance` will then pick the appropriate tier (scalar / AVX / AVX+FMA / AVX2+FMA / AVX-512) based on the host. From Python, use `lance.simd_info()` to verify which tier was selected. + ## Sample Workflow 1. Fork the repo diff --git a/python/python/lance/__init__.py b/python/python/lance/__init__.py index 61d94b5550a..3479a177824 100644 --- a/python/python/lance/__init__.py +++ b/python/python/lance/__init__.py @@ -47,6 +47,7 @@ ScanStatistics, bytes_read_counter, iops_counter, + simd_info, ) from .mem_wal import ( ExecutionPlan, @@ -115,6 +116,7 @@ "json_to_schema", "schema_to_json", "set_logger", + "simd_info", "write_dataset", "FFILanceTableProvider", "IndexProgress", diff --git a/python/python/tests/test_lance.py b/python/python/tests/test_lance.py index 0162e370665..2ad3af2e029 100644 --- a/python/python/tests/test_lance.py +++ b/python/python/tests/test_lance.py @@ -248,6 +248,28 @@ def test_io_counters(tmp_path): assert lance.bytes_read_counter() > starting_bytes +def test_simd_info(): + info = lance.simd_info() + assert info["tier"] in ( + "none", + "sse", + "avx", + "avx_fma", + "avx2", + "avx512", + "avx512_fp16", + "neon", + "lsx", + "lasx", + ) + assert isinstance(info["target_arch"], str) and info["target_arch"] + if info["target_arch"] == "x86_64": + # The x86_64 ABI mandates SSE2. + assert "sse2" in info["host_features"] + else: + assert info["host_features"] == [] + + @pytest.mark.parametrize( "row_param, column_name", [("with_row_id", "_rowid"), ("with_row_address", "_rowaddr")], diff --git a/python/src/lib.rs b/python/src/lib.rs index 466d4ea90f2..d40784643ab 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -322,6 +322,7 @@ fn lance(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(language_model_home))?; m.add_wrapped(wrap_pyfunction!(bytes_read_counter))?; m.add_wrapped(wrap_pyfunction!(iops_counter))?; + m.add_wrapped(wrap_pyfunction!(simd_info))?; m.add_wrapped(wrap_pyfunction!(stable_version))?; // Debug functions m.add_wrapped(wrap_pyfunction!(debug::format_schema))?; @@ -340,6 +341,29 @@ fn iops_counter() -> PyResult { Ok(::lance::io::iops_counter()) } +/// Returns a dict describing which SIMD tier the lance runtime dispatches to +/// on this host, plus the raw CPU feature flags it detected. +/// +/// Mirrors `pyarrow.runtime_info()`: a cheap, transparent way to verify that +/// the host is hitting the expected SIMD tier (e.g., `"avx512_fp16"`, +/// `"avx2"`) when debugging vector-search performance. +/// +/// Returns: +/// { +/// "tier": str, # e.g. "avx2", "avx_fma", "neon", "none" +/// "target_arch": str, # e.g. "x86_64", "aarch64", "loongarch64" +/// "host_features": list[str], # raw CPU feature flags (x86_64 only) +/// } +#[pyfunction] +pub fn simd_info(py: Python<'_>) -> PyResult> { + let info = lance_core::utils::cpu::simd_info(); + let dict = pyo3::types::PyDict::new(py); + dict.set_item("tier", info.tier.to_string())?; + dict.set_item("target_arch", info.target_arch)?; + dict.set_item("host_features", info.host_features)?; + Ok(dict.into()) +} + #[pyfunction(name = "bytes_read_counter")] fn bytes_read_counter() -> PyResult { Ok(::lance::io::bytes_read_counter()) diff --git a/rust/lance-core/src/utils/cpu.rs b/rust/lance-core/src/utils/cpu.rs index 4e7ab01871d..d3b200f671f 100644 --- a/rust/lance-core/src/utils/cpu.rs +++ b/rust/lance-core/src/utils/cpu.rs @@ -1,14 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use std::fmt; use std::sync::LazyLock; -/// A level of SIMD support for some feature +/// A level of SIMD support for some feature. +/// +/// `#[non_exhaustive]` so future tiers (e.g. AVX-512 BF16, AMX) can be added +/// without breaking external `match` consumers. #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] pub enum SimdSupport { None, Neon, Sse, + /// AVX (256-bit float ops) but no FMA and no AVX2. + /// Intel Sandy Bridge / Ivy Bridge. + Avx, + /// AVX + FMA but no AVX2. + /// AMD Piledriver / Steamroller / FX-7500. + AvxFma, Avx2, Avx512, Avx512FP16, @@ -16,6 +27,122 @@ pub enum SimdSupport { Lasx, } +impl fmt::Display for SimdSupport { + /// Formats the tier name in lowercase, matching pyarrow's + /// `runtime_info().simd_level` convention. + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let name = match self { + Self::None => "none", + Self::Neon => "neon", + Self::Sse => "sse", + Self::Avx => "avx", + Self::AvxFma => "avx_fma", + Self::Avx2 => "avx2", + Self::Avx512 => "avx512", + Self::Avx512FP16 => "avx512_fp16", + Self::Lsx => "lsx", + Self::Lasx => "lasx", + }; + f.write_str(name) + } +} + +/// Snapshot of the SIMD tier lance dispatches to on the current host, plus the +/// raw CPU features detected for diagnostic purposes. +/// +/// Mirrors the role of `pyarrow.runtime_info()`: a single, cheap call users can +/// make to verify which SIMD tier the runtime selected and what underlying +/// features the host advertises. +#[derive(Debug, Clone)] +pub struct SimdInfo { + /// The SIMD tier lance dispatches to at runtime on this host. + pub tier: SimdSupport, + /// The architecture name (e.g. "x86_64", "aarch64", "loongarch64"). + pub target_arch: &'static str, + /// Raw CPU feature flags detected on this host (x86_64 only; empty on + /// other architectures). Each entry is a feature name like "avx2", + /// "fma", "avx512f", "popcnt", etc. + pub host_features: Vec<&'static str>, +} + +/// Returns a snapshot of the SIMD tier lance is using on this host along with +/// the raw CPU feature flags that drove the decision. +/// +/// Useful for performance debugging and giving users a way to verify which +/// dispatch tier they are hitting without rebuilding lance. +pub fn simd_info() -> SimdInfo { + SimdInfo { + tier: *SIMD_SUPPORT, + target_arch: std::env::consts::ARCH, + host_features: detect_host_features(), + } +} + +#[cfg(target_arch = "x86_64")] +fn detect_host_features() -> Vec<&'static str> { + // Each call must be inline: `is_x86_feature_detected!` does its own custom + // input parsing and rejects feature names received via a `macro_rules!` + // `:literal` metavariable on some toolchains. + let mut features = Vec::new(); + if is_x86_feature_detected!("sse2") { + features.push("sse2"); + } + if is_x86_feature_detected!("sse3") { + features.push("sse3"); + } + if is_x86_feature_detected!("ssse3") { + features.push("ssse3"); + } + if is_x86_feature_detected!("sse4.1") { + features.push("sse4.1"); + } + if is_x86_feature_detected!("sse4.2") { + features.push("sse4.2"); + } + if is_x86_feature_detected!("popcnt") { + features.push("popcnt"); + } + if is_x86_feature_detected!("avx") { + features.push("avx"); + } + if is_x86_feature_detected!("avx2") { + features.push("avx2"); + } + if is_x86_feature_detected!("fma") { + features.push("fma"); + } + if is_x86_feature_detected!("f16c") { + features.push("f16c"); + } + if is_x86_feature_detected!("bmi1") { + features.push("bmi1"); + } + if is_x86_feature_detected!("bmi2") { + features.push("bmi2"); + } + if is_x86_feature_detected!("avx512f") { + features.push("avx512f"); + } + if is_x86_feature_detected!("avx512bw") { + features.push("avx512bw"); + } + if is_x86_feature_detected!("avx512cd") { + features.push("avx512cd"); + } + if is_x86_feature_detected!("avx512dq") { + features.push("avx512dq"); + } + if is_x86_feature_detected!("avx512vl") { + features.push("avx512vl"); + } + features +} + +#[cfg(not(target_arch = "x86_64"))] +fn detect_host_features() -> Vec<&'static str> { + Vec::new() +} + /// Support for SIMD operations pub static SIMD_SUPPORT: LazyLock = LazyLock::new(|| { #[cfg(all(target_arch = "aarch64", any(target_os = "ios", target_os = "tvos")))] @@ -44,6 +171,12 @@ pub static SIMD_SUPPORT: LazyLock = LazyLock::new(|| { } } else if is_x86_feature_detected!("avx2") { SimdSupport::Avx2 + } else if is_x86_feature_detected!("avx") && is_x86_feature_detected!("fma") { + // AMD Piledriver / Steamroller / FX-7500: 256-bit float ops + FMA but no AVX2. + SimdSupport::AvxFma + } else if is_x86_feature_detected!("avx") { + // Intel Sandy Bridge / Ivy Bridge: 256-bit float ops without FMA. + SimdSupport::Avx } else { SimdSupport::None } @@ -138,3 +271,46 @@ mod aarch64 { false } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn simd_info_exposes_tier() { + let info = simd_info(); + assert_eq!(info.target_arch, std::env::consts::ARCH); + // Tier should match the detected SIMD support. + assert_eq!(info.tier, *SIMD_SUPPORT); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn simd_info_features_include_baseline() { + let info = simd_info(); + // The x86_64 ABI mandates SSE2, so it must always be present on this + // architecture. + assert!(info.host_features.contains(&"sse2")); + } + + #[cfg(not(target_arch = "x86_64"))] + #[test] + fn simd_info_features_empty_off_x86_64() { + let info = simd_info(); + assert!(info.host_features.is_empty()); + } + + #[test] + fn simd_support_display_matches_lowercase_convention() { + assert_eq!(SimdSupport::Avx2.to_string(), "avx2"); + assert_eq!(SimdSupport::AvxFma.to_string(), "avx_fma"); + assert_eq!(SimdSupport::Avx.to_string(), "avx"); + assert_eq!(SimdSupport::Avx512.to_string(), "avx512"); + assert_eq!(SimdSupport::Avx512FP16.to_string(), "avx512_fp16"); + assert_eq!(SimdSupport::None.to_string(), "none"); + assert_eq!(SimdSupport::Neon.to_string(), "neon"); + assert_eq!(SimdSupport::Sse.to_string(), "sse"); + assert_eq!(SimdSupport::Lsx.to_string(), "lsx"); + assert_eq!(SimdSupport::Lasx.to_string(), "lasx"); + } +} diff --git a/rust/lance-linalg/proptest-regressions/distance/cosine.txt b/rust/lance-linalg/proptest-regressions/distance/cosine.txt new file mode 100644 index 00000000000..04e6dcb967f --- /dev/null +++ b/rust/lance-linalg/proptest-regressions/distance/cosine.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 679c764b02e2726ec2d943b3abdc7d7d3565e168d4081b155860c1e69b616fce # shrinks to (x, y) = ([-1.2933521e-15, -1.1194652e-8, 7.447341e-32, -0.0, -14247490.0, -8.55e-43, -4576.872, 0.009181444], [4.926e-42, -9.267065e-15, -0.0, 7e-45, -0.00036999694, -0.0, -5.7810876e-8, 1.9273644e-21]) diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index 995191b77eb..259be19f602 100644 --- a/rust/lance-linalg/src/distance/cosine.rs +++ b/rust/lance-linalg/src/distance/cosine.rs @@ -17,12 +17,12 @@ use arrow_array::{ use arrow_schema::DataType; use half::{bf16, f16}; use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray}; -use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] -use lance_core::utils::cpu::SimdSupport; +#[allow(unused_imports)] +use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use super::{Dot, norm_l2::norm_l2}; use super::{Normalize, dot::dot}; +#[allow(unused_imports)] use crate::simd::{ FloatSimd, SIMD, f32::{f32x8, f32x16}, @@ -127,6 +127,10 @@ impl Cosine for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::cosine_bf16_lsx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels in `bf16_kernel::*` are compiled with + // `-march=haswell` minimum (which requires AVX2), so they cannot + // run on AVX-only or AVX+FMA hosts. Scalar is the correct route. _ => cosine_scalar(x, x_norm, y), } } @@ -179,86 +183,203 @@ impl Cosine for f16 { SimdSupport::Lsx => unsafe { kernel::cosine_f16_lsx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => cosine_scalar(x, x_norm, y), } } } -/// f32 kernels for Cosine +/// f32 single-vector cosine helpers used by `cosine_batch` for fixed +/// dimensions 8 and 16. +/// +/// These were previously a single generic `cosine_once` but the +/// monomorphizations have to dispatch on `SIMD_SUPPORT` for the SIMD path +/// to stay correct under any compile baseline. Splitting them into two +/// concrete entry points keeps the dispatch site flat and lets each width +/// route to a `#[target_feature]` AVX2 inner function. mod f32 { use super::*; - // TODO: how can we explicitly infer N? #[inline] - pub(super) fn cosine_once, const N: usize>( - x: &[f32], - x_norm: f32, - y: &[f32], - ) -> f32 { - let x = unsafe { S::load_unaligned(x.as_ptr()) }; - let y = unsafe { S::load_unaligned(y.as_ptr()) }; - let y2 = y * y; - let xy = x * y; - 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + pub(super) fn cosine_once_8(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + cosine_once_x86::cosine_once_8_avx512(x, x_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + cosine_once_x86::cosine_once_8_avx_fma(x, x_norm, y) + }, + SimdSupport::Avx => unsafe { cosine_once_x86::cosine_once_8_avx(x, x_norm, y) }, + _ => cosine_once_8_scalar(x, x_norm, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_once_8_other(x, x_norm, y) + } } -} -impl Cosine for f32 { #[inline] - fn cosine_fast(x: &[Self], x_norm: Self, other: &[Self]) -> f32 { - let dim = x.len(); - let unrolled_len = dim / 16 * 16; - let mut y_norm16 = f32x16::zeros(); - let mut xy16 = f32x16::zeros(); - for i in (0..unrolled_len).step_by(16) { - unsafe { - let x = f32x16::load_unaligned(x.as_ptr().add(i)); - let y = f32x16::load_unaligned(other.as_ptr().add(i)); - xy16.multiply_add(x, y); - y_norm16.multiply_add(y, y); + pub(super) fn cosine_once_16(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + cosine_once_x86::cosine_once_16_avx512(x, x_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + cosine_once_x86::cosine_once_16_avx_fma(x, x_norm, y) + }, + SimdSupport::Avx => unsafe { cosine_once_x86::cosine_once_16_avx(x, x_norm, y) }, + _ => cosine_once_16_scalar(x, x_norm, y), } } - let aligned_len = dim / 8 * 8; - let mut y_norm8 = f32x8::zeros(); - let mut xy8 = f32x8::zeros(); - for i in (unrolled_len..aligned_len).step_by(8) { - unsafe { - let x = f32x8::load_unaligned(x.as_ptr().add(i)); - let y = f32x8::load_unaligned(other.as_ptr().add(i)); - xy8.multiply_add(x, y); - y_norm8.multiply_add(y, y); - } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_once_16_other(x, x_norm, y) } - let y_norm = - y_norm16.reduce_sum() + y_norm8.reduce_sum() + norm_l2(&other[aligned_len..]).powi(2); - let xy = - xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &other[aligned_len..]); - 1.0 - xy / x_norm / y_norm.sqrt() } + /// Portable scalar `cosine_once` for length-8 vectors. Matches the SIMD + /// path modulo summation order. + #[cfg(target_arch = "x86_64")] #[inline] - fn cosine_with_norms(x: &[Self], x_norm: Self, y_norm: Self, y: &[Self]) -> Self { - let dim = x.len(); - let unrolled_len = dim / 16 * 16; - let mut xy16 = f32x16::zeros(); - for i in (0..unrolled_len).step_by(16) { - unsafe { - let x = f32x16::load_unaligned(x.as_ptr().add(i)); - let y = f32x16::load_unaligned(y.as_ptr().add(i)); - xy16.multiply_add(x, y); - } + pub(super) fn cosine_once_8_scalar(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let mut xy = 0.0f32; + let mut y2 = 0.0f32; + for i in 0..8 { + xy += x[i] * y[i]; + y2 += y[i] * y[i]; } - let aligned_len = dim / 8 * 8; - let mut xy8 = f32x8::zeros(); - for i in (unrolled_len..aligned_len).step_by(8) { - unsafe { - let x = f32x8::load_unaligned(x.as_ptr().add(i)); - let y = f32x8::load_unaligned(y.as_ptr().add(i)); - xy8.multiply_add(x, y); - } + 1.0 - xy / x_norm / y2.sqrt() + } + + #[cfg(target_arch = "x86_64")] + #[inline] + pub(super) fn cosine_once_16_scalar(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let mut xy = 0.0f32; + let mut y2 = 0.0f32; + for i in 0..16 { + xy += x[i] * y[i]; + y2 += y[i] * y[i]; } - let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &y[aligned_len..]); - 1.0 - xy / x_norm / y_norm + 1.0 - xy / x_norm / y2.sqrt() + } + + #[cfg(target_arch = "x86_64")] + pub(super) mod cosine_once_x86 { + use std::arch::x86_64::*; + + use super::{f32x8, f32x16}; + use crate::simd::SIMD; + + /// AVX + FMA path for 8-lane cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_once_8_avx_fma(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x8::load_unaligned(x.as_ptr()); + let yv = f32x8::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX + FMA path for 16-lane cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_once_16_avx_fma(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x16::load_unaligned(x.as_ptr()); + let yv = f32x16::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX-only path for 8-lane cosine (no FMA): body unchanged from AVX2 path; gated on Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_once_8_avx(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x8::load_unaligned(x.as_ptr()); + let yv = f32x8::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX-only path for 16-lane cosine (no FMA): body unchanged from AVX2 path; gated on Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_once_16_avx(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = f32x16::load_unaligned(x.as_ptr()); + let yv = f32x16::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + /// AVX-512 path for 8-lane cosine: masked load into a `__m512` lower half, reduce. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_once_8_avx512(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + // mask 0x00FF: load the lower 8 f32 lanes, zero the upper 8. + let mask: __mmask16 = 0x00FF; + let xv = _mm512_maskz_loadu_ps(mask, x.as_ptr()); + let yv = _mm512_maskz_loadu_ps(mask, y.as_ptr()); + let xy = _mm512_mul_ps(xv, yv); + let y2 = _mm512_mul_ps(yv, yv); + let xy_sum = _mm512_reduce_add_ps(xy); + let y2_sum = _mm512_reduce_add_ps(y2); + 1.0 - xy_sum / x_norm / y2_sum.sqrt() + } + + /// AVX-512 path for 16-lane cosine: single full-width `__m512` load (16 f32 fits one `zmm`). + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_once_16_avx512(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = _mm512_loadu_ps(x.as_ptr()); + let yv = _mm512_loadu_ps(y.as_ptr()); + let xy = _mm512_mul_ps(xv, yv); + let y2 = _mm512_mul_ps(yv, yv); + let xy_sum = _mm512_reduce_add_ps(xy); + let y2_sum = _mm512_reduce_add_ps(y2); + 1.0 - xy_sum / x_norm / y2_sum.sqrt() + } + } + + #[cfg(not(target_arch = "x86_64"))] + #[inline] + fn cosine_once_8_other(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = unsafe { f32x8::load_unaligned(x.as_ptr()) }; + let yv = unsafe { f32x8::load_unaligned(y.as_ptr()) }; + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + + #[cfg(not(target_arch = "x86_64"))] + #[inline] + fn cosine_once_16_other(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + let xv = unsafe { f32x16::load_unaligned(x.as_ptr()) }; + let yv = unsafe { f32x16::load_unaligned(y.as_ptr()) }; + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } +} + +impl Cosine for f32 { + #[inline] + fn cosine_fast(x: &[Self], x_norm: Self, other: &[Self]) -> f32 { + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 inner kernel on capable hosts, or a portable scalar fallback. + cosine_fast_f32_dispatched(x, x_norm, other) + } + + #[inline] + fn cosine_with_norms(x: &[Self], x_norm: Self, y_norm: Self, y: &[Self]) -> Self { + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 inner kernel on capable hosts, or a portable scalar fallback. + cosine_with_norms_f32_dispatched(x, x_norm, y_norm, y) } fn cosine_batch<'a>( @@ -272,12 +393,12 @@ impl Cosine for f32 { 8 => Box::new( batch .chunks_exact(dimension) - .map(move |y| f32::cosine_once::(x, x_norm, y)), + .map(move |y| f32::cosine_once_8(x, x_norm, y)), ), 16 => Box::new( batch .chunks_exact(dimension) - .map(move |y| f32::cosine_once::(x, x_norm, y)), + .map(move |y| f32::cosine_once_16(x, x_norm, y)), ), _ => Box::new( batch @@ -291,34 +412,101 @@ impl Cosine for f32 { impl Cosine for f64 { #[inline] fn cosine_fast(x: &[Self], x_norm: f32, y: &[Self]) -> f32 { - use crate::simd::f64::{f64x4, f64x8}; - use crate::simd::{FloatSimd, SIMD}; + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 inner kernel on capable hosts, or a portable scalar fallback. + cosine_fast_f64_dispatched(x, x_norm, y) + } +} + +/// Fast cosine for f64, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the SIMD +/// primitives in `crate::simd::f64`. +#[inline] +fn cosine_fast_f64_dispatched(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + f64_x86::cosine_fast_avx512(x, x_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + f64_x86::cosine_fast_avx_fma(x, x_norm, y) + }, + SimdSupport::Avx => unsafe { f64_x86::cosine_fast_avx(x, x_norm, y) }, + _ => cosine_scalar(x, x_norm, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_fast_f64_simd_other(x, x_norm, y) + } +} + +/// AVX2 + FMA implementation of the f64 cosine_fast kernel. +/// +/// Lives in a `#[target_feature]`-annotated function so the SIMD primitives +/// in `crate::simd::f64` (which use raw AVX intrinsics) inline correctly +/// even when the compile baseline does not have AVX2 enabled. Caller must +/// ensure the host supports AVX2 + FMA. +#[cfg(target_arch = "x86_64")] +mod f64_x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX-512 path for f64 fast cosine: 8-wide `__m512d` xy/yy with `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_fast_avx512(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc_xy = _mm512_setzero_pd(); + let mut acc_yy = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let xv = _mm512_loadu_pd(x.as_ptr().add(i)); + let yv = _mm512_loadu_pd(y.as_ptr().add(i)); + acc_xy = _mm512_fmadd_pd(xv, yv, acc_xy); + acc_yy = _mm512_fmadd_pd(yv, yv, acc_yy); + } + + let mut xy = _mm512_reduce_add_pd(acc_xy); + let mut yy = _mm512_reduce_add_pd(acc_yy); + for i in unrolled_len..dim { + xy += x[i] * y[i]; + yy += y[i] * y[i]; + } + + let y_norm_sq = yy as f32; + let xy_f32 = xy as f32; + 1.0 - xy_f32 / x_norm / y_norm_sq.sqrt() + } + /// AVX + FMA path for f64 fast cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_fast_avx_fma(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { let dim = x.len(); let unrolled_len = dim / 8 * 8; let mut y_norm8 = f64x8::zeros(); let mut xy8 = f64x8::zeros(); for i in (0..unrolled_len).step_by(8) { - unsafe { - let xv = f64x8::load_unaligned(x.as_ptr().add(i)); - let yv = f64x8::load_unaligned(y.as_ptr().add(i)); - xy8.multiply_add(xv, yv); - y_norm8.multiply_add(yv, yv); - } + let xv = f64x8::load_unaligned(x.as_ptr().add(i)); + let yv = f64x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); } let aligned_len = dim / 4 * 4; let mut y_norm4 = f64x4::zeros(); let mut xy4 = f64x4::zeros(); for i in (unrolled_len..aligned_len).step_by(4) { - unsafe { - let xv = f64x4::load_unaligned(x.as_ptr().add(i)); - let yv = f64x4::load_unaligned(y.as_ptr().add(i)); - xy4.multiply_add(xv, yv); - y_norm4.multiply_add(yv, yv); - } + let xv = f64x4::load_unaligned(x.as_ptr().add(i)); + let yv = f64x4::load_unaligned(y.as_ptr().add(i)); + xy4.multiply_add(xv, yv); + y_norm4.multiply_add(yv, yv); } - let tail_y_norm: Self = y[aligned_len..].iter().map(|&v| v * v).sum(); - let tail_xy: Self = x[aligned_len..] + let tail_y_norm: f64 = y[aligned_len..].iter().map(|&v| v * v).sum(); + let tail_xy: f64 = x[aligned_len..] .iter() .zip(y[aligned_len..].iter()) .map(|(&a, &b)| a * b) @@ -328,6 +516,364 @@ impl Cosine for f64 { let xy = (xy8.reduce_sum() + xy4.reduce_sum() + tail_xy) as f32; 1.0 - xy / x_norm / y_norm_sq.sqrt() } + + /// AVX-only path for f64 fast cosine (no FMA): `_mm256_mul_pd` + `_mm256_add_pd` per iteration; tail handled inline. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_fast_avx(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + let dim = x.len(); + let aligned_len = dim / 4 * 4; + + let mut acc_xy = _mm256_setzero_pd(); + let mut acc_yy = _mm256_setzero_pd(); + for i in (0..aligned_len).step_by(4) { + let xv = _mm256_loadu_pd(x.as_ptr().add(i)); + let yv = _mm256_loadu_pd(y.as_ptr().add(i)); + acc_xy = _mm256_add_pd(acc_xy, _mm256_mul_pd(xv, yv)); + acc_yy = _mm256_add_pd(acc_yy, _mm256_mul_pd(yv, yv)); + } + + let xy_main = hsum256_pd_local(acc_xy); + let yy_main = hsum256_pd_local(acc_yy); + + let tail_y_norm: f64 = y[aligned_len..].iter().map(|&v| v * v).sum(); + let tail_xy: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + let y_norm_sq = (yy_main + tail_y_norm) as f32; + let xy = (xy_main + tail_xy) as f32; + 1.0 - xy / x_norm / y_norm_sq.sqrt() + } + + /// Horizontal sum of an `__m256d` register. Local to this `f64_x86` + /// mod for cache locality with its inner kernels; matches the + /// per-mod hsum-helper pattern already used in `dot.rs`/`norm_l2.rs`. + #[inline] + #[target_feature(enable = "avx")] + unsafe fn hsum256_pd_local(v: __m256d) -> f64 { + let lo = _mm256_castpd256_pd128(v); + let hi = _mm256_extractf128_pd(v, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + _mm_cvtsd_f64(sum64) + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn cosine_fast_f64_simd_other(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::{FloatSimd, SIMD}; + + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + let mut y_norm8 = f64x8::zeros(); + let mut xy8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + unsafe { + let xv = f64x8::load_unaligned(x.as_ptr().add(i)); + let yv = f64x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); + } + } + let aligned_len = dim / 4 * 4; + let mut y_norm4 = f64x4::zeros(); + let mut xy4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + unsafe { + let xv = f64x4::load_unaligned(x.as_ptr().add(i)); + let yv = f64x4::load_unaligned(y.as_ptr().add(i)); + xy4.multiply_add(xv, yv); + y_norm4.multiply_add(yv, yv); + } + } + let tail_y_norm: f64 = y[aligned_len..].iter().map(|&v| v * v).sum(); + let tail_xy: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + let y_norm_sq = (y_norm8.reduce_sum() + y_norm4.reduce_sum() + tail_y_norm) as f32; + let xy = (xy8.reduce_sum() + xy4.reduce_sum() + tail_xy) as f32; + 1.0 - xy / x_norm / y_norm_sq.sqrt() +} + +/// Cosine for f32 with known norms, runtime-dispatched via `SIMD_SUPPORT` +/// on x86_64 (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses +/// the auto-vectorised scalar loop. +#[inline] +fn cosine_with_norms_f32_dispatched(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + f32_x86::cosine_with_norms_avx512(x, x_norm, y_norm, y) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + f32_x86::cosine_with_norms_avx_fma(x, x_norm, y_norm, y) + }, + SimdSupport::Avx => unsafe { f32_x86::cosine_with_norms_avx(x, x_norm, y_norm, y) }, + _ => cosine_scalar_fast(x, x_norm, y, y_norm), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_with_norms_f32_simd_other(x, x_norm, y_norm, y) + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn cosine_with_norms_f32_simd_other(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + unsafe { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(y.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + } + } + let aligned_len = dim / 8 * 8; + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + unsafe { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + } + } + let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &y[aligned_len..]); + 1.0 - xy / x_norm / y_norm +} + +/// Fast cosine for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// `simd::f32` primitives, unconditionally backed by NEON / LSX. +#[inline] +fn cosine_fast_f32_dispatched(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + f32_x86::cosine_fast_avx512(x, x_norm, other) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { + f32_x86::cosine_fast_avx_fma(x, x_norm, other) + }, + SimdSupport::Avx => unsafe { f32_x86::cosine_fast_avx(x, x_norm, other) }, + _ => cosine_scalar(x, x_norm, other), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_fast_f32_simd_other(x, x_norm, other) + } +} + +/// AVX2 + FMA implementation of the f32 fast cosine kernel. +/// +/// Lives in a `#[target_feature]`-annotated function so the SIMD primitives +/// in `crate::simd::f32` (which use raw AVX intrinsics) inline correctly +/// even when the compile baseline does not have AVX2 enabled. Caller must +/// ensure the host supports AVX2 + FMA. +#[cfg(target_arch = "x86_64")] +mod f32_x86 { + use std::arch::x86_64::*; + + use super::{dot, f32x8, f32x16, norm_l2}; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX + FMA path for f32 fast cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_fast_avx_fma(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut y_norm16 = f32x16::zeros(); + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(other.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + y_norm16.multiply_add(yv, yv); + } + let aligned_len = dim / 8 * 8; + let mut y_norm8 = f32x8::zeros(); + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(other.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); + } + let y_norm = + y_norm16.reduce_sum() + y_norm8.reduce_sum() + norm_l2(&other[aligned_len..]).powi(2); + let xy = + xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() + } + + /// AVX-only path for f32 fast cosine (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration; tail via trait-routed `dot`/`norm_l2`. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_fast_avx(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let aligned_len = dim / 8 * 8; + + let mut acc_xy = _mm256_setzero_ps(); + let mut acc_yy = _mm256_setzero_ps(); + for i in (0..aligned_len).step_by(8) { + let xv = _mm256_loadu_ps(x.as_ptr().add(i)); + let yv = _mm256_loadu_ps(other.as_ptr().add(i)); + acc_xy = _mm256_add_ps(acc_xy, _mm256_mul_ps(xv, yv)); + acc_yy = _mm256_add_ps(acc_yy, _mm256_mul_ps(yv, yv)); + } + + let xy_main = hsum256_ps_local(acc_xy); + let yy_main = hsum256_ps_local(acc_yy); + + let y_norm = yy_main + norm_l2(&other[aligned_len..]).powi(2); + let xy = xy_main + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() + } + + /// Horizontal sum of an `__m256` register. Local to this `f32_x86` mod + /// because the sibling `mod x86` in `norm_l2.rs` keeps its own copy + /// for cache locality with its inner kernels; duplicating the helper + /// here avoids a cross-module unsafe re-export and matches the pattern + /// already used in `dot.rs::x86`. + #[inline] + #[target_feature(enable = "avx")] + unsafe fn hsum256_ps_local(v: __m256) -> f32 { + let lo = _mm256_castps256_ps128(v); + let hi = _mm256_extractf128_ps(v, 1); + let sum128 = _mm_add_ps(lo, hi); + let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); + _mm_cvtss_f32(sum32) + } + + /// AVX-512 path for f32 fast cosine: 16-wide `__m512` xy/yy with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_fast_avx512(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc_xy = _mm512_setzero_ps(); + let mut acc_yy = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let xv = _mm512_loadu_ps(x.as_ptr().add(i)); + let yv = _mm512_loadu_ps(other.as_ptr().add(i)); + acc_xy = _mm512_fmadd_ps(xv, yv, acc_xy); + acc_yy = _mm512_fmadd_ps(yv, yv, acc_yy); + } + + let mut xy = _mm512_reduce_add_ps(acc_xy); + let mut yy = _mm512_reduce_add_ps(acc_yy); + for i in unrolled_len..dim { + xy += x[i] * other[i]; + yy += other[i] * other[i]; + } + + 1.0 - xy / x_norm / yy.sqrt() + } + + /// AVX-512 path for f32 cosine with known norms: 16-wide `__m512` with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn cosine_with_norms_avx512(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let xv = _mm512_loadu_ps(x.as_ptr().add(i)); + let yv = _mm512_loadu_ps(y.as_ptr().add(i)); + acc = _mm512_fmadd_ps(xv, yv, acc); + } + + let mut xy = _mm512_reduce_add_ps(acc); + for i in unrolled_len..dim { + xy += x[i] * y[i]; + } + + 1.0 - xy / x_norm / y_norm + } + + /// AVX + FMA path for f32 cosine with known norms. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn cosine_with_norms_avx_fma(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(y.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + } + let aligned_len = dim / 8 * 8; + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(y.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + } + let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &y[aligned_len..]); + 1.0 - xy / x_norm / y_norm + } + + /// AVX-only path for f32 cosine with known norms (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration; tail via trait-routed `dot`. + #[target_feature(enable = "avx")] + pub unsafe fn cosine_with_norms_avx(x: &[f32], x_norm: f32, y_norm: f32, y: &[f32]) -> f32 { + let dim = x.len(); + let aligned_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..aligned_len).step_by(8) { + let xv = _mm256_loadu_ps(x.as_ptr().add(i)); + let yv = _mm256_loadu_ps(y.as_ptr().add(i)); + acc = _mm256_add_ps(acc, _mm256_mul_ps(xv, yv)); + } + + let xy_main = hsum256_ps_local(acc); + let xy = xy_main + dot(&x[aligned_len..], &y[aligned_len..]); + 1.0 - xy / x_norm / y_norm + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn cosine_fast_f32_simd_other(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut y_norm16 = f32x16::zeros(); + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + unsafe { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(other.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + y_norm16.multiply_add(yv, yv); + } + } + let aligned_len = dim / 8 * 8; + let mut y_norm8 = f32x8::zeros(); + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + unsafe { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(other.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); + } + } + let y_norm = + y_norm16.reduce_sum() + y_norm8.reduce_sum() + norm_l2(&other[aligned_len..]).powi(2); + let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() } /// Fallback non-SIMD implementation @@ -437,6 +983,25 @@ pub fn cosine_distance_arrow_batch( } } +/// Portable scalar reference cosine over f64 inputs. Used by parity tests +/// to compare against every dispatched per-tier inner kernel. Computes +/// `1 - xy / (x_norm * y_norm_sq.sqrt())` in f64 then casts to f32, matching +/// the reduction order of the dispatched kernels. +#[cfg(test)] +fn cosine_fast_scalar(x: &[f64], x_norm: f32, y: &[f64]) -> f32 { + let xy: f64 = x.iter().zip(y.iter()).map(|(&a, &b)| a * b).sum(); + let y_norm_sq: f64 = y.iter().map(|&v| v * v).sum(); + 1.0 - (xy as f32) / x_norm / (y_norm_sq as f32).sqrt() +} + +/// Portable scalar reference cosine when both norms are known. Mirrors +/// `cosine_with_norms_f32_dispatched` for parity testing. +#[cfg(test)] +fn cosine_with_norms_scalar(x: &[f64], x_norm: f32, y_norm: f32, y: &[f64]) -> f32 { + let xy: f64 = x.iter().zip(y.iter()).map(|(&a, &b)| a * b).sum(); + 1.0 - (xy as f32) / x_norm / y_norm +} + #[cfg(test)] mod tests { use super::*; @@ -559,5 +1124,388 @@ mod tests { prop_assume!(norm_l2(&y) > 1e-20); do_cosine_test(&x, &y)?; } + + /// Cross-backend parity for the f32 cosine_fast kernel. Exercises the + /// scalar fallback (`cosine_scalar`) against the dispatched SIMD path + /// so the runtime fallback is exercised even on AVX2-capable CI hosts. + #[test] + fn test_cosine_fast_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + let simd = ::cosine_fast(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity for the f32 cosine_fast kernel. Early-returns + /// on hosts without AVX-512F so the test stays portable. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = cosine_scalar(&x, x_norm, &y); + let avx512 = unsafe { f32_x86::cosine_fast_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the f32 cosine_fast kernel. Covers + /// the AMD Piledriver / Steamroller / FX-7500 tier. Early-returns + /// on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = cosine_scalar(&x, x_norm, &y); + let avx_fma = unsafe { f32_x86::cosine_fast_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the f32 cosine_fast kernel. Covers + /// the Intel Sandy Bridge / Ivy Bridge tier. Early-returns on + /// hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = cosine_scalar(&x, x_norm, &y); + let avx = unsafe { f32_x86::cosine_fast_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// Cross-backend parity for the f32 cosine_with_norms kernel. + /// Exercises the scalar fallback (`cosine_scalar_fast`) against the + /// dispatched SIMD path so the runtime fallback is exercised even on + /// AVX2-capable CI hosts. + #[test] + fn test_cosine_with_norms_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_with_norms_scalar(&x_f64, x_norm, y_norm, &y_f64); + let simd = ::cosine_with_norms(&x, x_norm, y_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity for the f32 cosine_with_norms kernel. + /// Early-returns on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_with_norms_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let scalar = cosine_scalar_fast(&x, x_norm, &y, y_norm); + let avx512 = unsafe { f32_x86::cosine_with_norms_avx512(&x, x_norm, y_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the f32 cosine_with_norms kernel. + /// Covers the AMD Piledriver / Steamroller / FX-7500 tier. + /// Early-returns on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_with_norms_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let scalar = cosine_scalar_fast(&x, x_norm, &y, y_norm); + let avx_fma = unsafe { f32_x86::cosine_with_norms_avx_fma(&x, x_norm, y_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the f32 cosine_with_norms kernel. + /// Covers the Intel Sandy Bridge / Ivy Bridge tier. Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_with_norms_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let y_norm = norm_l2(&y); + let scalar = cosine_scalar_fast(&x, x_norm, &y, y_norm); + let avx = unsafe { f32_x86::cosine_with_norms_avx(&x, x_norm, y_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// Cross-backend parity for the f64 cosine_fast kernel. Uses the + /// hand-rolled `cosine_fast_scalar` (not the trait-routed + /// `cosine_scalar`, which would itself dispatch through `dot::`) + /// so the reference stays free of any AVX path on AVX2-capable hosts. + #[test] + fn test_cosine_fast_f64_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let simd = ::cosine_fast(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity for the f64 cosine_fast kernel. Early-returns + /// on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f64_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let avx512 = unsafe { f64_x86::cosine_fast_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the f64 cosine_fast kernel. Covers + /// the AMD Piledriver / Steamroller / FX-7500 tier. Early-returns + /// on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f64_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let avx_fma = unsafe { f64_x86::cosine_fast_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the f64 cosine_fast kernel. Covers + /// the Intel Sandy Bridge / Ivy Bridge tier. Early-returns on + /// hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_fast_f64_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-20); + prop_assume!(norm_l2(&y) > 1e-20); + let x_norm = norm_l2(&x); + let scalar = cosine_fast_scalar(&x, x_norm, &y); + let avx = unsafe { f64_x86::cosine_fast_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// Parity check for `cosine_once_8` (despecialised 8-lane width). + /// + /// The `epsilon = 1e-6` clause handles the case where the proptest + /// generator produces inputs with extreme dynamic range (e.g., mixing + /// `1e-43` with `1e7` in the same vector). When the dot product is + /// dominated by one large term and the cosine result is near zero, + /// the f32-precision SIMD path and the f64-precision scalar reference + /// can legitimately differ by more than `max_relative = 1e-3` of the + /// (near-zero) result. The absolute epsilon catches these without + /// masking real bugs (where the absolute error would be > 1e-6). + #[test] + fn test_cosine_once_8_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + let simd = f32::cosine_once_8(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3, epsilon = 1e-6)); + } + + /// Parity check for `cosine_once_16` (despecialised 16-lane width). + /// See `test_cosine_once_8_scalar_simd_parity` for `epsilon` rationale. + #[test] + fn test_cosine_once_16_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + let simd = f32::cosine_once_16(&x, x_norm, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3, epsilon = 1e-6)); + } + + /// AVX-512-direct parity for the 8-lane cosine_once kernel. Verifies + /// the masked-load (mask 0x00FF) AVX-512 implementation produces + /// the same result as the scalar reference. Early-returns on hosts + /// without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_8_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_8_scalar(&x, x_norm, &y); + let avx512 = + unsafe { super::f32::cosine_once_x86::cosine_once_8_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX-512-direct parity for the 16-lane cosine_once kernel. Verifies + /// the full-width `__m512` load implementation produces the same + /// result as the scalar reference. Early-returns on hosts without + /// AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_16_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_16_scalar(&x, x_norm, &y); + let avx512 = + unsafe { super::f32::cosine_once_x86::cosine_once_16_avx512(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the 8-lane cosine_once kernel. + /// Covers the AMD Piledriver / Steamroller / FX-7500 tier. + /// Early-returns on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_8_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_8_scalar(&x, x_norm, &y); + let avx_fma = + unsafe { super::f32::cosine_once_x86::cosine_once_8_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX + FMA-direct parity for the 16-lane cosine_once kernel. + /// Covers the AMD Piledriver / Steamroller / FX-7500 tier. + /// Early-returns on hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_16_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_16_scalar(&x, x_norm, &y); + let avx_fma = + unsafe { super::f32::cosine_once_x86::cosine_once_16_avx_fma(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the 8-lane cosine_once kernel. + /// Covers the Intel Sandy Bridge / Ivy Bridge tier. Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_8_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 8..9) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_8_scalar(&x, x_norm, &y); + let avx = unsafe { super::f32::cosine_once_x86::cosine_once_8_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } + + /// AVX-only-direct parity for the 16-lane cosine_once kernel. + /// Covers the Intel Sandy Bridge / Ivy Bridge tier. Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_once_16_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 16..17) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + prop_assume!(norm_l2(&x) > 1e-10); + prop_assume!(norm_l2(&y) > 1e-10); + let x_norm = norm_l2(&x); + let scalar = super::f32::cosine_once_16_scalar(&x, x_norm, &y); + let avx = unsafe { super::f32::cosine_once_x86::cosine_once_16_avx(&x, x_norm, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); + } } } diff --git a/rust/lance-linalg/src/distance/cosine_u8.rs b/rust/lance-linalg/src/distance/cosine_u8.rs index b2d06d35c31..59833e020a7 100644 --- a/rust/lance-linalg/src/distance/cosine_u8.rs +++ b/rust/lance-linalg/src/distance/cosine_u8.rs @@ -199,6 +199,10 @@ fn select_backend() -> CosineU8AccumFn { if is_x86_feature_detected!("avx2") { return |a, b| unsafe { x86::cosine_u8_accum_avx2(a, b) }; } + // AvxFma and Avx hosts (AMD Piledriver / Steamroller, Intel Sandy + // Bridge / Ivy Bridge) fall through to scalar: the AVX2 inner uses + // `vpmaddubsw` / `vpmaddwd` integer ops which neither AVX nor + // AVX+FMA provides. } cosine_u8_accum_scalar diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 5903d24e0e5..7d3e7a36ff0 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -14,9 +14,8 @@ use arrow_schema::DataType; use half::{bf16, f16}; use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray}; use lance_core::assume_eq; -use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] -use lance_core::utils::cpu::SimdSupport; +#[allow(unused_imports)] +use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use num_traits::{AsPrimitive, Num, real::Real}; use crate::Result; @@ -161,6 +160,9 @@ impl Dot for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::dot_bf16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => dot_scalar::(x, y), } } @@ -214,6 +216,9 @@ impl Dot for f16 { SimdSupport::Lsx => unsafe { kernel::dot_f16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => dot_scalar::(x, y), } } @@ -222,10 +227,43 @@ impl Dot for f16 { impl Dot for f32 { #[inline] fn dot(x: &[Self], y: &[Self]) -> f32 { - dot_scalar::(x, y) + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 or AVX-512 inner kernel on capable hosts, or a portable + // scalar fallback. Same shape as the f64 sibling and the existing + // u8 distance kernels in `dot_u8.rs`. + dot_f32_dispatched(x, y) } } +/// Dot product for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// auto-vectorised scalar loop. +#[inline] +fn dot_f32_dispatched(x: &[f32], y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::dot_f32_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::dot_f32_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::dot_f32_avx(x, y) }, + _ => dot_f32_scalar(x, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + dot_f32_scalar(x, y) + } +} + +/// Portable scalar dot product for f32. Used as the x86_64 fallback when no +/// AVX2 is detected, and as the only path on non-x86 architectures. The +/// `LANES = 16` chunking matches the explicit-SIMD inner kernels above. +#[inline] +fn dot_f32_scalar(x: &[f32], y: &[f32]) -> f32 { + dot_scalar::(x, y) +} + impl Dot for f64 { #[inline] fn dot(x: &[Self], y: &[Self]) -> f32 { @@ -233,9 +271,207 @@ impl Dot for f64 { } } -/// Explicit SIMD dot product for f64. +/// Dot product for f64, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the SIMD +/// primitives in `crate::simd::f64`, unconditionally backed by NEON / LSX-LASX. #[inline] fn dot_f64_simd(x: &[f64], y: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::dot_f64_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::dot_f64_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::dot_f64_avx(x, y) }, + _ => dot_f64_scalar(x, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + dot_f64_simd_other(x, y) + } +} + +/// Portable scalar dot product for f64. Used as the x86_64 fallback when no +/// AVX2 is detected, and exposed for cross-backend parity testing. +#[cfg(target_arch = "x86_64")] +#[inline] +fn dot_f64_scalar(x: &[f64], y: &[f64]) -> f32 { + x.iter().zip(y.iter()).map(|(&a, &b)| a * b).sum::() as f32 +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX-512 path for f64: 8-wide `__m512d` with `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn dot_f64_avx512(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm512_loadu_pd(x.as_ptr().add(i)); + let b = _mm512_loadu_pd(y.as_ptr().add(i)); + acc = _mm512_fmadd_pd(a, b, acc); + } + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + (_mm512_reduce_add_pd(acc) + tail) as f32 + } + + /// AVX + FMA path for f64. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn dot_f64_avx_fma(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + let a = f64x8::load_unaligned(x.as_ptr().add(i)); + let b = f64x8::load_unaligned(y.as_ptr().add(i)); + acc8.multiply_add(a, b); + } + + let aligned_len = dim / 4 * 4; + let mut acc4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + let a = f64x4::load_unaligned(x.as_ptr().add(i)); + let b = f64x4::load_unaligned(y.as_ptr().add(i)); + acc4.multiply_add(a, b); + } + + let tail: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + (acc8.reduce_sum() + acc4.reduce_sum() + tail) as f32 + } + + /// AVX-only path for f64 (no FMA): `_mm256_mul_pd` + `_mm256_add_pd` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn dot_f64_avx(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 4 * 4; + + let mut acc = _mm256_setzero_pd(); + for i in (0..unrolled_len).step_by(4) { + let a = _mm256_loadu_pd(x.as_ptr().add(i)); + let b = _mm256_loadu_pd(y.as_ptr().add(i)); + acc = _mm256_add_pd(acc, _mm256_mul_pd(a, b)); + } + + // Horizontal sum of __m256d -> f64. Two pairwise adds across lanes. + let lo = _mm256_castpd256_pd128(acc); + let hi = _mm256_extractf128_pd(acc, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + let acc_sum = _mm_cvtsd_f64(sum64); + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + (acc_sum + tail) as f32 + } + + /// Horizontal sum of an `__m256` register. Folds the upper 128-bit lane + /// into the lower, then sums lanes pairwise. Same shape as the helper in + /// the sibling `norm_l2.rs` mod x86; kept local rather than hoisted to a + /// shared module to avoid a one-helper module file. + #[inline] + #[target_feature(enable = "avx")] + unsafe fn hsum256_ps(v: __m256) -> f32 { + let lo = _mm256_castps256_ps128(v); + let hi = _mm256_extractf128_ps(v, 1); + let sum128 = _mm_add_ps(lo, hi); + let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); + _mm_cvtss_f32(sum32) + } + + /// AVX-512 path for f32: 16-wide `__m512` with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn dot_f32_avx512(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let a = _mm512_loadu_ps(x.as_ptr().add(i)); + let b = _mm512_loadu_ps(y.as_ptr().add(i)); + acc = _mm512_fmadd_ps(a, b, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + _mm512_reduce_add_ps(acc) + tail + } + + /// AVX + FMA path for f32. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn dot_f32_avx_fma(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + acc = _mm256_fmadd_ps(a, b, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + hsum256_ps(acc) + tail + } + + /// AVX-only path for f32 (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn dot_f32_avx(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + acc = _mm256_add_ps(acc, _mm256_mul_ps(a, b)); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| a * b) + .sum(); + + hsum256_ps(acc) + tail + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn dot_f64_simd_other(x: &[f64], y: &[f64]) -> f32 { use crate::simd::f64::{f64x4, f64x8}; use crate::simd::{FloatSimd, SIMD}; @@ -480,5 +716,151 @@ mod tests { fn test_dot_f64((x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048)){ do_dot_test(&x, &y)?; } + + /// Cross-backend parity: scalar fallback must match the dispatched + /// SIMD path within numerical tolerance. Exercises `dot_f64_scalar` + /// directly so the runtime fallback is exercised even on AVX2-capable + /// CI hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + let scalar = dot_f64_scalar(&x, &y); + let simd = dot_f64_simd(&x, &y); + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, simd, epsilon = max_error)); + } + + /// Parity check for `dot_f32_dispatched` (Branch B exclusive: the + /// auto-vectorised scalar dot path). The dispatched kernel must + /// agree with a portable f64-precision scalar reference within + /// numerical tolerance. The reference is hand-rolled here to keep + /// this test architecture-agnostic (the x86_64-only `dot_f64_scalar` + /// helper is gated above). + #[test] + fn test_dot_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let scalar = x_f64 + .iter() + .zip(y_f64.iter()) + .map(|(&a, &b)| a * b) + .sum::() as f32; + let simd = ::dot(&x, &y); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, simd, epsilon = max_error)); + } + + /// AVX-512-direct parity for f32: explicitly compares the scalar + /// fallback against the native f32 AVX-512 inner kernel on + /// AVX-512F-capable hosts. Early-returns on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = dot_f32_scalar(&x, &y); + let avx512 = unsafe { x86::dot_f32_avx512(&x, &y) }; + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, avx512, epsilon = max_error)); + } + + /// AVX + FMA-direct parity for the f32 dot kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = dot_f32_scalar(&x, &y); + let avx_fma = unsafe { x86::dot_f32_avx_fma(&x, &y) }; + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, avx_fma, epsilon = max_error)); + } + + /// AVX-only-direct parity for the f32 dot kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier. Early-returns on hosts without + /// AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = dot_f32_scalar(&x, &y); + let avx = unsafe { x86::dot_f32_avx(&x, &y) }; + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + let y_f64: Vec = y.iter().map(|&v| v as f64).collect(); + let max_error = max_error::(&x_f64, &y_f64); + prop_assert!(approx::relative_eq!(scalar, avx, epsilon = max_error)); + } + + /// AVX-512-direct parity: explicitly compares the scalar fallback + /// against the native AVX-512 inner kernel on AVX-512F-capable hosts + /// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4). Early-returns on + /// hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = dot_f64_scalar(&x, &y); + let avx512 = unsafe { x86::dot_f64_avx512(&x, &y) }; + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, avx512, epsilon = max_error)); + } + + /// AVX + FMA-direct parity for the f64 dot kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = dot_f64_scalar(&x, &y); + let avx_fma = unsafe { x86::dot_f64_avx_fma(&x, &y) }; + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, avx_fma, epsilon = max_error)); + } + + /// AVX-only-direct parity for the f64 dot kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier (AVX without FMA). Early-returns + /// on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_f64_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = dot_f64_scalar(&x, &y); + let avx = unsafe { x86::dot_f64_avx(&x, &y) }; + let max_error = max_error::(&x, &y); + prop_assert!(approx::relative_eq!(scalar, avx, epsilon = max_error)); + } } } diff --git a/rust/lance-linalg/src/distance/dot_u8.rs b/rust/lance-linalg/src/distance/dot_u8.rs index de5522cddfe..7b2e335f094 100644 --- a/rust/lance-linalg/src/distance/dot_u8.rs +++ b/rust/lance-linalg/src/distance/dot_u8.rs @@ -134,6 +134,10 @@ fn select_backend() -> DotU8Fn { if is_x86_feature_detected!("avx2") { return |a, b| unsafe { x86::dot_u8_avx2(a, b) }; } + // AvxFma and Avx hosts (AMD Piledriver / Steamroller, Intel Sandy + // Bridge / Ivy Bridge) fall through to scalar: the AVX2 inner uses + // `vpmaddubsw` / `vpmaddwd` integer ops which neither AVX nor + // AVX+FMA provides. } dot_u8_scalar diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index c47aedd749f..78886860401 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -20,7 +20,6 @@ use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray}; use lance_core::assume_eq; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] use lance_core::utils::cpu::SimdSupport; use num_traits::{AsPrimitive, Num}; @@ -191,6 +190,9 @@ impl L2 for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::l2_bf16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => l2_scalar::(x, y), } } @@ -244,6 +246,9 @@ impl L2 for f16 { SimdSupport::Lsx => unsafe { kernel::l2_f16_lsx(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => l2_scalar::(x, y), } } @@ -252,10 +257,42 @@ impl L2 for f16 { impl L2 for f32 { #[inline] fn l2(x: &[Self], y: &[Self]) -> f32 { - // 16 = 512 (avx512) / 8 bits / 4 (sizeof(f32)) - // See https://github.com/lance-format/lance/pull/2450. - l2_scalar::(x, y) + // Trait methods cannot carry `#[target_feature]` attributes, so the body + // lives in a free function that runtime-dispatches via `*SIMD_SUPPORT` + // to an AVX2 or AVX-512 inner kernel on capable hosts, or a portable + // scalar fallback. + l2_f32_dispatched(x, y) + } +} + +/// L2 distance for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// auto-vectorised scalar loop. +#[inline] +fn l2_f32_dispatched(x: &[f32], y: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::l2_f32_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::l2_f32_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::l2_f32_avx(x, y) }, + _ => l2_f32_scalar(x, y), + } } + #[cfg(not(target_arch = "x86_64"))] + { + l2_f32_scalar(x, y) + } +} + +/// Portable scalar L2 distance for f32. Used as the x86_64 fallback when no +/// AVX2 is detected, and as the only path on non-x86 architectures. The +/// `LANES = 16` chunking matches the explicit-SIMD inner kernels above. +#[inline] +fn l2_f32_scalar(x: &[f32], y: &[f32]) -> f32 { + // 16 = 512 (avx512) / 8 bits / 4 (sizeof(f32)) + // See https://github.com/lance-format/lance/pull/2450. + l2_scalar::(x, y) } impl L2 for f64 { @@ -265,9 +302,238 @@ impl L2 for f64 { } } -/// Explicit SIMD L2 distance for f64. +/// L2 distance for f64, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the SIMD +/// primitives in `crate::simd::f64`, unconditionally backed by NEON / LSX-LASX. #[inline] fn l2_f64_simd(x: &[f64], y: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { x86::l2_f64_avx512(x, y) }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::l2_f64_avx_fma(x, y) }, + SimdSupport::Avx => unsafe { x86::l2_f64_avx(x, y) }, + _ => l2_f64_scalar(x, y), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + l2_f64_simd_other(x, y) + } +} + +/// Portable scalar L2 distance for f64. Used as the x86_64 fallback when no +/// AVX2 is detected, and exposed for cross-backend parity testing. +#[cfg(target_arch = "x86_64")] +#[inline] +fn l2_f64_scalar(x: &[f64], y: &[f64]) -> f32 { + x.iter() + .zip(y.iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum::() as f32 +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX-512 path for f64: 8-wide `__m512d` with `vsubpd` + `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn l2_f64_avx512(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm512_loadu_pd(x.as_ptr().add(i)); + let b = _mm512_loadu_pd(y.as_ptr().add(i)); + let diff = _mm512_sub_pd(a, b); + acc = _mm512_fmadd_pd(diff, diff, acc); + } + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + (_mm512_reduce_add_pd(acc) + tail) as f32 + } + + /// AVX + FMA path for f64. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn l2_f64_avx_fma(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + let a = f64x8::load_unaligned(x.as_ptr().add(i)); + let b = f64x8::load_unaligned(y.as_ptr().add(i)); + let diff = a - b; + acc8.multiply_add(diff, diff); + } + + let aligned_len = dim / 4 * 4; + let mut acc4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + let a = f64x4::load_unaligned(x.as_ptr().add(i)); + let b = f64x4::load_unaligned(y.as_ptr().add(i)); + let diff = a - b; + acc4.multiply_add(diff, diff); + } + + let tail: f64 = x[aligned_len..] + .iter() + .zip(y[aligned_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + (acc8.reduce_sum() + acc4.reduce_sum() + tail) as f32 + } + + /// AVX-only path for f64 (no FMA): squared diff via `_mm256_mul_pd` + `_mm256_add_pd` for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn l2_f64_avx(x: &[f64], y: &[f64]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 4 * 4; + + let mut acc = _mm256_setzero_pd(); + for i in (0..unrolled_len).step_by(4) { + let a = _mm256_loadu_pd(x.as_ptr().add(i)); + let b = _mm256_loadu_pd(y.as_ptr().add(i)); + let diff = _mm256_sub_pd(a, b); + acc = _mm256_add_pd(acc, _mm256_mul_pd(diff, diff)); + } + + // Horizontal sum of __m256d -> f64. + let lo = _mm256_castpd256_pd128(acc); + let hi = _mm256_extractf128_pd(acc, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + let acc_sum = _mm_cvtsd_f64(sum64); + + let tail: f64 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + (acc_sum + tail) as f32 + } + + /// Horizontal sum of an `__m256` register. Folds the upper 128-bit lane + /// into the lower, then sums lanes pairwise. Same shape as the helper in + /// the sibling `norm_l2.rs` mod x86; kept local rather than hoisted to a + /// shared module to avoid a one-helper module file. + #[inline] + #[target_feature(enable = "avx")] + unsafe fn hsum256_ps(v: __m256) -> f32 { + let lo = _mm256_castps256_ps128(v); + let hi = _mm256_extractf128_ps(v, 1); + let sum128 = _mm_add_ps(lo, hi); + let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); + _mm_cvtss_f32(sum32) + } + + /// AVX-512 path for f32: 16-wide `__m512` with `vsubps` + `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn l2_f32_avx512(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let a = _mm512_loadu_ps(x.as_ptr().add(i)); + let b = _mm512_loadu_ps(y.as_ptr().add(i)); + let diff = _mm512_sub_ps(a, b); + acc = _mm512_fmadd_ps(diff, diff, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + _mm512_reduce_add_ps(acc) + tail + } + + /// AVX + FMA path for f32. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn l2_f32_avx_fma(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + let diff = _mm256_sub_ps(a, b); + acc = _mm256_fmadd_ps(diff, diff, acc); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + hsum256_ps(acc) + tail + } + + /// AVX-only path for f32 (no FMA): squared diff via `_mm256_mul_ps` + `_mm256_add_ps` for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn l2_f32_avx(x: &[f32], y: &[f32]) -> f32 { + let dim = x.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let a = _mm256_loadu_ps(x.as_ptr().add(i)); + let b = _mm256_loadu_ps(y.as_ptr().add(i)); + let diff = _mm256_sub_ps(a, b); + acc = _mm256_add_ps(acc, _mm256_mul_ps(diff, diff)); + } + + let tail: f32 = x[unrolled_len..] + .iter() + .zip(y[unrolled_len..].iter()) + .map(|(&a, &b)| { + let diff = a - b; + diff * diff + }) + .sum(); + + hsum256_ps(acc) + tail + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn l2_f64_simd_other(x: &[f64], y: &[f64]) -> f32 { use crate::simd::f64::{f64x4, f64x8}; use crate::simd::{FloatSimd, SIMD}; @@ -671,6 +937,136 @@ mod tests { fn test_l2_distance_f64((x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048)){ do_l2_test(&x, &y)?; } + + /// Cross-backend parity: scalar fallback must match the dispatched + /// SIMD path within numerical tolerance. Exercises `l2_f64_scalar` + /// directly so the runtime fallback is exercised even on AVX2-capable + /// CI hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + let scalar = l2_f64_scalar(&x, &y); + let simd = l2_f64_simd(&x, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-6)); + } + + /// Parity check for `l2_f32_dispatched` (Branch B exclusive: the + /// auto-vectorised scalar L2 path). The dispatched kernel must + /// agree with a portable f64-precision scalar reference within + /// numerical tolerance. The reference is hand-rolled here to keep + /// this test architecture-agnostic (the x86_64-only `l2_f64_scalar` + /// helper is gated above). + #[test] + fn test_l2_f32_scalar_simd_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + let scalar = x + .iter() + .zip(y.iter()) + .map(|(&a, &b)| ((a as f64) - (b as f64)).powi(2)) + .sum::() as f32; + let simd = ::l2(&x, &y); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity: explicitly compares the scalar fallback + /// against the native AVX-512 inner kernel on AVX-512F-capable hosts + /// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4). Early-returns on + /// hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = l2_f64_scalar(&x, &y); + let avx512 = unsafe { x86::l2_f64_avx512(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-6)); + } + + /// AVX + FMA-direct parity for the f64 L2 kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = l2_f64_scalar(&x, &y); + let avx_fma = unsafe { x86::l2_f64_avx_fma(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-6)); + } + + /// AVX-only-direct parity for the f64 L2 kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier. Early-returns on hosts without + /// AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f64_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f64, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = l2_f64_scalar(&x, &y); + let avx = unsafe { x86::l2_f64_avx(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-6)); + } + + /// AVX-512-direct parity for f32: explicitly compares the scalar + /// fallback against the native f32 AVX-512 inner kernel on + /// AVX-512F-capable hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f32_scalar_vs_avx512_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = l2_f32_scalar(&x, &y); + let avx512 = unsafe { x86::l2_f32_avx512(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-3)); + } + + /// AVX + FMA-direct parity for the f32 L2 kernel. Covers the AMD + /// Piledriver / Steamroller / FX-7500 tier. Early-returns on hosts + /// without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f32_scalar_vs_avx_fma_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = l2_f32_scalar(&x, &y); + let avx_fma = unsafe { x86::l2_f32_avx_fma(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-3)); + } + + /// AVX-only-direct parity for the f32 L2 kernel. Covers the Intel + /// Sandy Bridge / Ivy Bridge tier. Early-returns on hosts without + /// AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_f32_scalar_vs_avx_parity( + (x, y) in arbitrary_vector_pair(arbitrary_f32, 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = l2_f32_scalar(&x, &y); + let avx = unsafe { x86::l2_f32_avx(&x, &y) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-3)); + } } #[test] diff --git a/rust/lance-linalg/src/distance/l2_u8.rs b/rust/lance-linalg/src/distance/l2_u8.rs index 1b111f91338..5fbf8a3af55 100644 --- a/rust/lance-linalg/src/distance/l2_u8.rs +++ b/rust/lance-linalg/src/distance/l2_u8.rs @@ -143,6 +143,10 @@ fn select_backend() -> L2U8Fn { if is_x86_feature_detected!("avx2") { return |a, b| unsafe { x86::l2_u8_avx2(a, b) }; } + // AvxFma and Avx hosts (AMD Piledriver / Steamroller, Intel Sandy + // Bridge / Ivy Bridge) fall through to scalar: the AVX2 inner uses + // AVX2 integer ops (`vpsubusb` / `vpmaddwd`) which neither AVX nor + // AVX+FMA provides. } l2_u8_scalar diff --git a/rust/lance-linalg/src/distance/norm_l2.rs b/rust/lance-linalg/src/distance/norm_l2.rs index b1daf85ab3b..fa89f2ce44d 100644 --- a/rust/lance-linalg/src/distance/norm_l2.rs +++ b/rust/lance-linalg/src/distance/norm_l2.rs @@ -9,9 +9,7 @@ use arrow_array::types::{Float16Type, Float32Type, Float64Type}; use arrow_schema::DataType; use half::{bf16, f16}; #[allow(unused_imports)] -use lance_core::utils::cpu::SIMD_SUPPORT; -#[cfg(feature = "fp16kernels")] -use lance_core::utils::cpu::SimdSupport; +use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use num_traits::{AsPrimitive, Float, Num}; /// L2 normalization @@ -75,6 +73,9 @@ impl Normalize for f16 { SimdSupport::Lsx => unsafe { kernel::norm_l2_f16_lsx(vector.as_ptr(), vector.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the f16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => norm_l2_impl::(vector), } } @@ -126,6 +127,9 @@ impl Normalize for bf16 { SimdSupport::Lsx => unsafe { bf16_kernel::norm_l2_bf16_lsx(vector.as_ptr(), vector.len() as u32) }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the bf16 C kernels are compiled with `-march=haswell` minimum + // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts. _ => norm_l2_impl::(vector), } } @@ -134,10 +138,40 @@ impl Normalize for bf16 { impl Normalize for f32 { #[inline] fn norm_l2(vector: &[Self]) -> f32 { - norm_l2_impl::(vector) + norm_l2_f32_dispatched(vector) } } +/// L2 norm for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 +/// (AVX-512 / AVX2+FMA / AVX+FMA / AVX / scalar). Non-x86 uses the +/// auto-vectorised scalar loop. +#[inline] +fn norm_l2_f32_dispatched(vector: &[f32]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + x86::norm_l2_f32_avx512(vector) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::norm_l2_f32_avx_fma(vector) }, + SimdSupport::Avx => unsafe { x86::norm_l2_f32_avx(vector) }, + _ => norm_l2_f32_scalar(vector), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + norm_l2_f32_scalar(vector) + } +} + +/// Portable scalar L2 norm for f32. Used as the x86_64 fallback when no +/// AVX2 is detected, and as the only path on non-x86 architectures. The +/// `LANES = 16` chunking matches the explicit-SIMD inner kernels above. +#[inline] +fn norm_l2_f32_scalar(vector: &[f32]) -> f32 { + norm_l2_impl::(vector) +} + impl Normalize for f64 { #[inline] fn norm_l2(vector: &[Self]) -> f32 { @@ -145,11 +179,178 @@ impl Normalize for f64 { } } -/// Explicit SIMD implementation of L2 norm for f64. +/// L2 norm for f64. Runtime-dispatched to the best available backend. /// -/// Two-level unrolling: f64x8 main loop, f64x4 remainder, scalar tail. +/// On x86_64, dispatches via `SIMD_SUPPORT` to a native AVX-512 inner kernel +/// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4), an AVX2 + FMA kernel +/// (Haswell+), an AVX + FMA kernel (AMD Piledriver / Steamroller), an +/// AVX-only kernel (Intel Sandy Bridge / Ivy Bridge), or a portable scalar +/// fallback. The per-tier inner functions each carry their own +/// `#[target_feature]` so they stay correct under any compile baseline. +/// On aarch64 and loongarch64, the SIMD primitives in `crate::simd::f64` +/// are unconditionally backed by NEON / LSX-LASX respectively, so no +/// runtime gate is required. #[inline] pub fn norm_l2_f64_simd(vector: &[f64]) -> f32 { + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => unsafe { + x86::norm_l2_f64_avx512(vector) + }, + SimdSupport::Avx2 | SimdSupport::AvxFma => unsafe { x86::norm_l2_f64_avx_fma(vector) }, + SimdSupport::Avx => unsafe { x86::norm_l2_f64_avx(vector) }, + _ => norm_l2_f64_scalar(vector), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + norm_l2_f64_simd_other(vector) + } +} + +/// Portable scalar L2 norm. Used as the x86_64 fallback when no AVX2 is +/// detected, and exposed for cross-backend parity testing. +#[cfg(target_arch = "x86_64")] +#[inline] +fn norm_l2_f64_scalar(vector: &[f64]) -> f32 { + vector.iter().map(|v| v * v).sum::().sqrt() as f32 +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use std::arch::x86_64::*; + + use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::{FloatSimd, SIMD}; + + /// AVX-512 path for f64: 8-wide `__m512d` with `vfmadd231pd` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn norm_l2_f64_avx512(vector: &[f64]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm512_setzero_pd(); + for i in (0..unrolled_len).step_by(8) { + let v = _mm512_loadu_pd(vector.as_ptr().add(i)); + acc = _mm512_fmadd_pd(v, v, acc); + } + + let tail: f64 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (_mm512_reduce_add_pd(acc) + tail).sqrt() as f32 + } + + /// AVX + FMA path for f64. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn norm_l2_f64_avx_fma(vector: &[f64]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc8 = f64x8::zeros(); + for i in (0..unrolled_len).step_by(8) { + let v = f64x8::load_unaligned(vector.as_ptr().add(i)); + acc8.multiply_add(v, v); + } + + let aligned_len = dim / 4 * 4; + let mut acc4 = f64x4::zeros(); + for i in (unrolled_len..aligned_len).step_by(4) { + let v = f64x4::load_unaligned(vector.as_ptr().add(i)); + acc4.multiply_add(v, v); + } + + let tail: f64 = vector[aligned_len..].iter().map(|&v| v * v).sum(); + (acc8.reduce_sum() + acc4.reduce_sum() + tail).sqrt() as f32 + } + + /// AVX-only path for f64 (no FMA): `_mm256_mul_pd` + `_mm256_add_pd` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn norm_l2_f64_avx(vector: &[f64]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 4 * 4; + + let mut acc = _mm256_setzero_pd(); + for i in (0..unrolled_len).step_by(4) { + let v = _mm256_loadu_pd(vector.as_ptr().add(i)); + acc = _mm256_add_pd(acc, _mm256_mul_pd(v, v)); + } + + // Horizontal sum of __m256d -> f64. Two pairwise adds across lanes. + let lo = _mm256_castpd256_pd128(acc); + let hi = _mm256_extractf128_pd(acc, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + let acc_sum = _mm_cvtsd_f64(sum64); + + let tail: f64 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (acc_sum + tail).sqrt() as f32 + } + + /// Horizontal sum of an `__m256` register. Folds the upper 128-bit lane + /// into the lower, then sums lanes pairwise via two `vhaddps`. + #[inline] + #[target_feature(enable = "avx")] + unsafe fn hsum256_ps(v: __m256) -> f32 { + let lo = _mm256_castps256_ps128(v); + let hi = _mm256_extractf128_ps(v, 1); + let sum128 = _mm_add_ps(lo, hi); + let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); + _mm_cvtss_f32(sum32) + } + + /// AVX-512 path for f32: 16-wide `__m512` with `vfmadd231ps` per iteration. + #[target_feature(enable = "avx512f")] + pub unsafe fn norm_l2_f32_avx512(vector: &[f32]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 16 * 16; + + let mut acc = _mm512_setzero_ps(); + for i in (0..unrolled_len).step_by(16) { + let v = _mm512_loadu_ps(vector.as_ptr().add(i)); + acc = _mm512_fmadd_ps(v, v, acc); + } + + let tail: f32 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (_mm512_reduce_add_ps(acc) + tail).sqrt() + } + + /// AVX + FMA path for f32. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). + #[target_feature(enable = "avx,fma")] + pub unsafe fn norm_l2_f32_avx_fma(vector: &[f32]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let v = _mm256_loadu_ps(vector.as_ptr().add(i)); + acc = _mm256_fmadd_ps(v, v, acc); + } + + let tail: f32 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (hsum256_ps(acc) + tail).sqrt() + } + + /// AVX-only path for f32 (no FMA): `_mm256_mul_ps` + `_mm256_add_ps` per iteration for Sandy/Ivy Bridge. + #[target_feature(enable = "avx")] + pub unsafe fn norm_l2_f32_avx(vector: &[f32]) -> f32 { + let dim = vector.len(); + let unrolled_len = dim / 8 * 8; + + let mut acc = _mm256_setzero_ps(); + for i in (0..unrolled_len).step_by(8) { + let v = _mm256_loadu_ps(vector.as_ptr().add(i)); + acc = _mm256_add_ps(acc, _mm256_mul_ps(v, v)); + } + + let tail: f32 = vector[unrolled_len..].iter().map(|&v| v * v).sum(); + (hsum256_ps(acc) + tail).sqrt() + } +} + +#[cfg(not(target_arch = "x86_64"))] +#[inline] +fn norm_l2_f64_simd_other(vector: &[f64]) -> f32 { use crate::simd::f64::{f64x4, f64x8}; use crate::simd::{FloatSimd, SIMD}; @@ -291,5 +492,133 @@ mod tests { fn test_l2_norm_f64(data in prop::collection::vec(arbitrary_f64(), 4..4048)){ do_norm_l2_test(&data)?; } + + /// Cross-backend parity: scalar fallback must match the dispatched + /// SIMD path within numerical tolerance. Exercises `norm_l2_f64_scalar` + /// directly so the runtime fallback is exercised even on AVX2-capable + /// CI hosts. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_simd_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + let scalar = norm_l2_f64_scalar(&data); + let simd = norm_l2_f64_simd(&data); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-6)); + } + + /// Parity check for `norm_l2_f32_dispatched` (Branch B exclusive: the + /// auto-vectorised scalar L2-norm path). The dispatched kernel must + /// agree with a portable f64-precision scalar reference within + /// numerical tolerance. The reference is hand-rolled here to keep this + /// test architecture-agnostic (the x86_64-only `norm_l2_f64_scalar` + /// helper is gated above). + #[test] + fn test_l2_norm_f32_scalar_simd_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + let scalar = data.iter().map(|&v| (v as f64).powi(2)).sum::().sqrt() as f32; + let simd = ::norm_l2(&data); + prop_assert!(approx::relative_eq!(scalar, simd, max_relative = 1e-3)); + } + + /// AVX-512-direct parity: explicitly compares the scalar fallback + /// against the native AVX-512 inner kernel on AVX-512F-capable hosts + /// (Skylake-X+, Ice Lake, Sapphire Rapids, Zen 4). Early-returns on + /// hosts without AVX-512F so the test stays portable; CI runners with + /// AVX-512F exercise the `_mm512_*` path. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_vs_avx512_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = norm_l2_f64_scalar(&data); + let avx512 = unsafe { x86::norm_l2_f64_avx512(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-6)); + } + + /// AVX + FMA-direct parity for the f64 L2-norm kernel. Covers the + /// AMD Piledriver / Steamroller / FX-7500 tier. Early-returns on + /// hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_vs_avx_fma_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = norm_l2_f64_scalar(&data); + let avx_fma = unsafe { x86::norm_l2_f64_avx_fma(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-6)); + } + + /// AVX-only-direct parity for the f64 L2-norm kernel. Covers the + /// Intel Sandy Bridge / Ivy Bridge tier (AVX without FMA). + /// Early-returns on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_norm_f64_scalar_vs_avx_parity( + data in prop::collection::vec(arbitrary_f64(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = norm_l2_f64_scalar(&data); + let avx = unsafe { x86::norm_l2_f64_avx(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-6)); + } + + /// AVX-512-direct parity for the f32 L2-norm kernel. Explicitly + /// compares the scalar fallback against the native AVX-512 inner + /// kernel on AVX-512F-capable hosts. Early-returns on hosts without + /// AVX-512F; CI runners with AVX-512F exercise the `_mm512_*` path. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_norm_l2_f32_scalar_vs_avx512_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx512f") { + return Ok(()); + } + let scalar = norm_l2_f32_scalar(&data); + let avx512 = unsafe { x86::norm_l2_f32_avx512(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx512, max_relative = 1e-3)); + } + + /// AVX + FMA-direct parity for the f32 L2-norm kernel. Covers the + /// AMD Piledriver / Steamroller / FX-7500 tier. Early-returns on + /// hosts without both AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_norm_l2_f32_scalar_vs_avx_fma_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return Ok(()); + } + let scalar = norm_l2_f32_scalar(&data); + let avx_fma = unsafe { x86::norm_l2_f32_avx_fma(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx_fma, max_relative = 1e-3)); + } + + /// AVX-only-direct parity for the f32 L2-norm kernel. Covers the + /// Intel Sandy Bridge / Ivy Bridge tier (AVX without FMA). + /// Early-returns on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_norm_l2_f32_scalar_vs_avx_parity( + data in prop::collection::vec(arbitrary_f32(), 4..4048) + ) { + if !std::is_x86_feature_detected!("avx") { + return Ok(()); + } + let scalar = norm_l2_f32_scalar(&data); + let avx = unsafe { x86::norm_l2_f32_avx(&data) }; + prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-3)); + } } } diff --git a/rust/lance-linalg/src/simd/dist_table.rs b/rust/lance-linalg/src/simd/dist_table.rs index 626c1581b15..e337953ec10 100644 --- a/rust/lance-linalg/src/simd/dist_table.rs +++ b/rust/lance-linalg/src/simd/dist_table.rs @@ -73,6 +73,10 @@ pub fn sum_4bit_dist_table( ) } }, + // SimdSupport::AvxFma and SimdSupport::Avx fall through here: + // the AVX2 inner uses `_mm256_shuffle_epi8` / `_mm256_and_si256` / + // `_mm256_srli_epi16` / `_mm256_add_epi16` integer ops which + // neither AVX nor AVX+FMA provides. Scalar is the correct route. _ => sum_4bit_dist_table_scalar(code_len, codes, dist_table, dists), } } diff --git a/rust/lance-linalg/src/simd/f32.rs b/rust/lance-linalg/src/simd/f32.rs index 78042997121..638850af172 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -46,14 +46,26 @@ impl std::fmt::Debug for f32x8 { } impl f32x8 { + /// Gather 8 f32 values from `slice` at the offsets in `indices`. + /// + /// On x86_64 this uses the AVX2 `vgatherdps` instruction when the host + /// supports it (gated at runtime via `is_x86_feature_detected!`). On + /// other architectures (and on x86_64 hosts without AVX2) the function + /// falls back to a per-index scalar load followed by `Self::from(&out)`, + /// which goes through `load_unaligned` (NEON / LASX / `_mm256_loadu_ps` + /// depending on platform). Per-tier macro stamping (e.g., the + /// `multiversion` crate) was considered but doesn't fit here: the function + /// returns `Self` and `_mm256_i32gather_ps::<4>` requires the const-generic + /// stride to be a compile-time literal — neither composes with the macro. #[inline] pub fn gather(slice: &[f32], indices: &[i32; 8]) -> Self { #[cfg(target_arch = "x86_64")] - unsafe { - use super::i32::i32x8; - - let idx = i32x8::from(indices); - Self(_mm256_i32gather_ps::<4>(slice.as_ptr(), idx.0)) + { + if is_x86_feature_detected!("avx2") { + unsafe { gather_avx2(slice, indices) } + } else { + gather_scalar_x86(slice, indices) + } } #[cfg(target_arch = "aarch64")] @@ -94,6 +106,37 @@ impl f32x8 { } } +/// AVX2 gather. Caller must ensure the host supports AVX2 (gated by +/// the `is_x86_feature_detected!("avx2")` check in `f32x8::gather`). +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn gather_avx2(slice: &[f32], indices: &[i32; 8]) -> f32x8 { + use super::i32::i32x8; + + let idx = i32x8::from(indices); + f32x8(_mm256_i32gather_ps::<4>(slice.as_ptr(), idx.0)) +} + +/// Portable scalar gather for x86_64 hosts without AVX2. +#[cfg(target_arch = "x86_64")] +#[inline] +fn gather_scalar_x86(slice: &[f32], indices: &[i32; 8]) -> f32x8 { + let ptr = slice.as_ptr(); + unsafe { + let values = [ + *ptr.add(indices[0] as usize), + *ptr.add(indices[1] as usize), + *ptr.add(indices[2] as usize), + *ptr.add(indices[3] as usize), + *ptr.add(indices[4] as usize), + *ptr.add(indices[5] as usize), + *ptr.add(indices[6] as usize), + *ptr.add(indices[7] as usize), + ]; + f32x8::load_unaligned(values.as_ptr()) + } +} + impl From<&[f32]> for f32x8 { fn from(value: &[f32]) -> Self { unsafe { Self::load_unaligned(value.as_ptr()) } @@ -439,15 +482,18 @@ impl Mul for f32x8 { } } -/// 16 of 32-bit `f32` values. Use 512-bit SIMD if possible. +/// 16 of 32-bit `f32` values. Stored as a pair of 256-bit AVX vectors on +/// x86_64. Originally there was a sibling AVX-512 variant gated on +/// `target_feature = "avx512f"`, but no project CI configuration enables +/// `+avx512f` globally (and one of the avx512 arms still contained `todo!()`), +/// so the variant was dead code. Removed in the runtime-SIMD-dispatch +/// retrofit; per-tier dispatch happens in the kernel functions in +/// `crate::distance::*` via `match *SIMD_SUPPORT` + per-tier +/// `#[target_feature(enable = "...")]` inner functions. #[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] +#[cfg(target_arch = "x86_64")] #[derive(Clone, Copy)] pub struct f32x16(__m256, __m256); -#[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] -#[derive(Clone, Copy)] -pub struct f32x16(__m512); /// 16 of 32-bit `f32` values. Use 512-bit SIMD if possible. #[allow(non_camel_case_types)] @@ -486,11 +532,7 @@ impl<'a> From<&'a [f32; 16]> for f32x16 { impl SIMD for f32x16 { #[inline] fn splat(val: f32) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_set1_ps(val)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_set1_ps(val), _mm256_set1_ps(val)) } @@ -514,11 +556,7 @@ impl SIMD for f32x16 { #[inline] fn zeros() -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_setzero_ps()) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_setzero_ps(), _mm256_setzero_ps()) } @@ -534,14 +572,10 @@ impl SIMD for f32x16 { #[inline] unsafe fn load(ptr: *const f32) -> Self { - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_load_ps(ptr), _mm256_load_ps(ptr.add(8))) } - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_load_ps(ptr)) - } #[cfg(target_arch = "aarch64")] { Self::load_unaligned(ptr) @@ -557,14 +591,10 @@ impl SIMD for f32x16 { #[inline] unsafe fn load_unaligned(ptr: *const f32) -> Self { - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_loadu_ps(ptr), _mm256_loadu_ps(ptr.add(8))) } - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_loadu_ps(ptr)) - } #[cfg(target_arch = "aarch64")] { Self(vld1q_f32_x4(ptr)) @@ -580,11 +610,7 @@ impl SIMD for f32x16 { #[inline] unsafe fn store(&self, ptr: *mut f32) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_store_ps(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_store_ps(ptr, self.0); _mm256_store_ps(ptr.add(8), self.1); @@ -602,11 +628,7 @@ impl SIMD for f32x16 { #[inline] unsafe fn store_unaligned(&self, ptr: *mut f32) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_storeu_ps(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_storeu_ps(ptr, self.0); _mm256_storeu_ps(ptr.add(8), self.1); @@ -622,12 +644,9 @@ impl SIMD for f32x16 { } } + #[inline] fn reduce_sum(&self) -> f32 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_add_ps(0xFFFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let mut sum = _mm256_add_ps(self.0, self.1); // Shift and add vector, until only 1 value left. @@ -657,11 +676,7 @@ impl SIMD for f32x16 { #[inline] fn reduce_min(&self) -> f32 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_min_ps(0xFFFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let mut m1 = _mm256_min_ps(self.0, self.1); let mut m2 = _mm256_permute2f128_ps(m1, m1, 1); @@ -695,11 +710,7 @@ impl SIMD for f32x16 { #[inline] fn min(&self, rhs: &Self) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_min_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_min_ps(self.0, rhs.0), _mm256_min_ps(self.1, rhs.1)) } @@ -718,19 +729,12 @@ impl SIMD for f32x16 { } } + #[inline] fn find(&self, val: f32) -> Option { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - // let tgt = _mm512_set1_ps(val); - // let mask = _mm512_cmpeq_ps_mask(self.0, tgt); - // if mask != 0 { - // return Some(mask.trailing_zeros() as i32); - // } - todo!() - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { - // _mm256_cmpeq_ps_mask requires "avx512l". + // _mm256_cmpeq_ps_mask requires AVX-512 (avx512f); use a scalar scan here + // since we only require AVX2. for i in 0..16 { if self.as_array().get_unchecked(i) == &val { return Some(i as i32); @@ -774,11 +778,7 @@ impl SIMD for f32x16 { impl FloatSimd for f32x16 { #[inline] fn multiply_add(&mut self, a: Self, b: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_fmadd_ps(a.0, b.0, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_fmadd_ps(a.0, b.0, self.0); self.1 = _mm256_fmadd_ps(a.1, b.1, self.1); @@ -803,11 +803,7 @@ impl Add for f32x16 { #[inline] fn add(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_add_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_add_ps(self.0, rhs.0), _mm256_add_ps(self.1, rhs.1)) } @@ -830,11 +826,7 @@ impl Add for f32x16 { impl AddAssign for f32x16 { #[inline] fn add_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_add_ps(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_add_ps(self.0, rhs.0); self.1 = _mm256_add_ps(self.1, rhs.1); @@ -859,11 +851,7 @@ impl Mul for f32x16 { #[inline] fn mul(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_mul_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_mul_ps(self.0, rhs.0), _mm256_mul_ps(self.1, rhs.1)) } @@ -888,11 +876,7 @@ impl Sub for f32x16 { #[inline] fn sub(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_sub_ps(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_sub_ps(self.0, rhs.0), _mm256_sub_ps(self.1, rhs.1)) } @@ -915,11 +899,7 @@ impl Sub for f32x16 { impl SubAssign for f32x16 { #[inline] fn sub_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_sub_ps(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_sub_ps(self.0, rhs.0); self.1 = _mm256_sub_ps(self.1, rhs.1); @@ -946,6 +926,12 @@ mod tests { #[test] fn test_basic_ops() { + // The `f32x8` constructor / load / store paths are AVX-only on x86_64 + // after the baseline lowering; skip on hosts that don't support AVX2. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = (0..8).map(|f| f as f32).collect::>(); let b = (10..18).map(|f| f as f32).collect::>(); @@ -983,6 +969,10 @@ mod tests { #[test] fn test_f32x8_cmp_ops() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = [1.0_f32, 2.0, 5.0, 6.0, 7.0, 3.0, 2.0, 1.0]; let b = [2.0_f32, 1.0, 4.0, 5.0, 9.0, 5.0, 6.0, 2.0]; let c = [2.0_f32, 1.0, 4.0, 5.0, 7.0, 3.0, 2.0, 1.0]; @@ -1007,6 +997,10 @@ mod tests { #[test] fn test_basic_f32x16_ops() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = (0..16).map(|f| f as f32).collect::>(); let b = (10..26).map(|f| f as f32).collect::>(); @@ -1041,6 +1035,10 @@ mod tests { #[test] fn test_f32x16_cmp_ops() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = [ 1.0_f32, 2.0, 5.0, 6.0, 7.0, 3.0, 2.0, 1.0, -0.5, 5.0, 6.0, 7.0, 8.0, 9.0, 1.0, 2.0, ]; @@ -1074,6 +1072,13 @@ mod tests { #[test] fn test_f32x8_gather() { + // `f32x8::gather` does its own runtime AVX2 detection, but the + // returned `f32x8` itself is `__m256`-backed on x86_64 and the + // `reduce_sum` helper requires AVX. Skip on pre-Haswell hosts. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = (0..256).map(|f| f as f32).collect::>(); let idx = [0_i32, 4, 8, 12, 16, 20, 24, 29]; let v = f32x8::gather(&a, &idx); diff --git a/rust/lance-linalg/src/simd/f64.rs b/rust/lance-linalg/src/simd/f64.rs index 32c0d389e5b..0274c8ad6f1 100644 --- a/rust/lance-linalg/src/simd/f64.rs +++ b/rust/lance-linalg/src/simd/f64.rs @@ -355,14 +355,15 @@ impl Mul for f64x4 { // f64x8: 8 × f64 values (512-bit SIMD or 2 × 256-bit) // --------------------------------------------------------------------------- -/// 8 of 64-bit `f64` values. Uses 512-bit SIMD if possible. +/// 8 of 64-bit `f64` values. Stored as a pair of 256-bit AVX vectors on +/// x86_64. Originally there was a sibling AVX-512 variant gated on +/// `target_feature = "avx512f"`, but no project CI configuration enables +/// `+avx512f` globally, so the variant was dead code. Removed in the +/// runtime-SIMD-dispatch retrofit; per-tier dispatch happens in the kernel +/// functions in `crate::distance::*` via `match *SIMD_SUPPORT` + per-tier +/// `#[target_feature(enable = "...")]` inner functions. #[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] -#[derive(Clone, Copy)] -pub struct f64x8(__m512d); - -#[allow(non_camel_case_types)] -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] +#[cfg(target_arch = "x86_64")] #[derive(Clone, Copy)] pub struct f64x8(__m256d, __m256d); @@ -401,11 +402,7 @@ impl<'a> From<&'a [f64; 8]> for f64x8 { impl SIMD for f64x8 { #[inline] fn splat(val: f64) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_set1_pd(val)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_set1_pd(val), _mm256_set1_pd(val)) } @@ -423,11 +420,7 @@ impl SIMD for f64x8 { #[inline] fn zeros() -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_setzero_pd()) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_setzero_pd(), _mm256_setzero_pd()) } @@ -443,11 +436,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn load(ptr: *const f64) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_load_pd(ptr)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_load_pd(ptr), _mm256_load_pd(ptr.add(4))) } @@ -466,11 +455,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn load_unaligned(ptr: *const f64) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_loadu_pd(ptr)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_loadu_pd(ptr), _mm256_loadu_pd(ptr.add(4))) } @@ -489,11 +474,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn store(&self, ptr: *mut f64) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_store_pd(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_store_pd(ptr, self.0); _mm256_store_pd(ptr.add(4), self.1); @@ -512,11 +493,7 @@ impl SIMD for f64x8 { #[inline] unsafe fn store_unaligned(&self, ptr: *mut f64) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_storeu_pd(ptr, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { _mm256_storeu_pd(ptr, self.0); _mm256_storeu_pd(ptr.add(4), self.1); @@ -533,12 +510,9 @@ impl SIMD for f64x8 { } } + #[inline] fn reduce_sum(&self) -> f64 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_add_pd(0xFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let sum = _mm256_add_pd(self.0, self.1); let hi = _mm256_permute2f128_pd(sum, sum, 1); @@ -561,11 +535,7 @@ impl SIMD for f64x8 { #[inline] fn reduce_min(&self) -> f64 { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - _mm512_mask_reduce_min_pd(0xFF, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { let m = _mm256_min_pd(self.0, self.1); let hi = _mm256_permute2f128_pd(m, m, 1); @@ -592,11 +562,7 @@ impl SIMD for f64x8 { #[inline] fn min(&self, rhs: &Self) -> Self { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_min_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_min_pd(self.0, rhs.0), _mm256_min_pd(self.1, rhs.1)) } @@ -613,6 +579,7 @@ impl SIMD for f64x8 { } } + #[inline] fn find(&self, val: f64) -> Option { unsafe { for i in 0..8 { @@ -628,11 +595,7 @@ impl SIMD for f64x8 { impl FloatSimd for f64x8 { #[inline] fn multiply_add(&mut self, a: Self, b: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_fmadd_pd(a.0, b.0, self.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_fmadd_pd(a.0, b.0, self.0); self.1 = _mm256_fmadd_pd(a.1, b.1, self.1); @@ -657,11 +620,7 @@ impl Add for f64x8 { #[inline] fn add(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_add_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_add_pd(self.0, rhs.0), _mm256_add_pd(self.1, rhs.1)) } @@ -682,11 +641,7 @@ impl Add for f64x8 { impl AddAssign for f64x8 { #[inline] fn add_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_add_pd(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_add_pd(self.0, rhs.0); self.1 = _mm256_add_pd(self.1, rhs.1); @@ -711,11 +666,7 @@ impl Mul for f64x8 { #[inline] fn mul(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_mul_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_mul_pd(self.0, rhs.0), _mm256_mul_pd(self.1, rhs.1)) } @@ -738,11 +689,7 @@ impl Sub for f64x8 { #[inline] fn sub(self, rhs: Self) -> Self::Output { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - Self(_mm512_sub_pd(self.0, rhs.0)) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { Self(_mm256_sub_pd(self.0, rhs.0), _mm256_sub_pd(self.1, rhs.1)) } @@ -763,11 +710,7 @@ impl Sub for f64x8 { impl SubAssign for f64x8 { #[inline] fn sub_assign(&mut self, rhs: Self) { - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - unsafe { - self.0 = _mm512_sub_pd(self.0, rhs.0) - } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx512f")))] + #[cfg(target_arch = "x86_64")] unsafe { self.0 = _mm256_sub_pd(self.0, rhs.0); self.1 = _mm256_sub_pd(self.1, rhs.1); @@ -793,6 +736,12 @@ mod tests { #[test] fn test_f64x4_basic_ops() { + // The `f64x4` constructor / load / store paths are AVX-only on x86_64 + // after the baseline lowering; skip on hosts that don't support AVX2. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = [1.0_f64, 2.0, 3.0, 4.0]; let b = [5.0_f64, 6.0, 7.0, 8.0]; @@ -814,6 +763,10 @@ mod tests { #[test] fn test_f64x4_fma() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = [1.0_f64, 2.0, 3.0, 4.0]; let b = [2.0_f64, 3.0, 4.0, 5.0]; @@ -826,6 +779,10 @@ mod tests { #[test] fn test_f64x4_min() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a = [1.0_f64, 5.0, 2.0, 8.0]; let b = [3.0_f64, 2.0, 4.0, 1.0]; let simd_a: f64x4 = (&a).into(); @@ -838,6 +795,10 @@ mod tests { #[test] fn test_f64x8_basic_ops() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a: [f64; 8] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; let b: [f64; 8] = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]; @@ -856,6 +817,10 @@ mod tests { #[test] fn test_f64x8_fma() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a: [f64; 8] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; let b: [f64; 8] = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]; @@ -869,6 +834,10 @@ mod tests { #[test] fn test_f64x8_min() { + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } let a: [f64; 8] = [5.0, 1.0, 8.0, 3.0, 9.0, 2.0, 7.0, 4.0]; let b: [f64; 8] = [2.0, 6.0, 3.0, 7.0, 1.0, 8.0, 4.0, 9.0]; let simd_a: f64x8 = (&a).into(); From 6141fed39b4d4e77dd2c6685ba05e9af2e0fc023 Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Wed, 10 Jun 2026 11:47:06 -0700 Subject: [PATCH 02/15] fix(lance-linalg): gate f32 cosine dispatch behind cfg(target_feature=avx2) On AVX2-baseline builds (the default haswell wheel), cosine_batch uses inlined non-target_feature kernels (base-equivalent) plus a single per-batch AVX-512 check, avoiding the per-vector runtime-dispatch + target_feature tax that regressed the modern path. Sub-AVX2 builds keep the runtime dispatch. Co-Authored-By: Claude Opus 4.8 --- rust/lance-linalg/src/distance/cosine.rs | 210 +++++++++++++++++++++++ 1 file changed, 210 insertions(+) diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index 259be19f602..c6615fa6757 100644 --- a/rust/lance-linalg/src/distance/cosine.rs +++ b/rust/lance-linalg/src/distance/cosine.rs @@ -363,6 +363,148 @@ mod f32 { let xy = xv * yv; 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() } + + /// Batch-level SIMD dispatch: the tier is chosen once by the caller, and the + /// whole `chunks_exact` loop runs inside one `#[target_feature]` context so + /// the per-vector `cosine_once_*` / `cosine_fast` kernels inline (no + /// per-vector `*SIMD_SUPPORT` branch, no per-vector call boundary). Used for + /// the AVX-512 path on the default wheel and for all tiers on sub-AVX2 builds. + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "avx,fma")] + pub(super) unsafe fn cosine_batch_avx_fma( + x: &[f32], + x_norm: f32, + batch: &[f32], + dimension: usize, + ) -> Vec { + match dimension { + 8 => batch + .chunks_exact(8) + .map(|y| unsafe { cosine_once_x86::cosine_once_8_avx_fma(x, x_norm, y) }) + .collect(), + 16 => batch + .chunks_exact(16) + .map(|y| unsafe { cosine_once_x86::cosine_once_16_avx_fma(x, x_norm, y) }) + .collect(), + _ => batch + .chunks_exact(dimension) + .map(|y| unsafe { super::f32_x86::cosine_fast_avx_fma(x, x_norm, y) }) + .collect(), + } + } + + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "avx512f")] + pub(super) unsafe fn cosine_batch_avx512( + x: &[f32], + x_norm: f32, + batch: &[f32], + dimension: usize, + ) -> Vec { + match dimension { + 8 => batch + .chunks_exact(8) + .map(|y| unsafe { cosine_once_x86::cosine_once_8_avx512(x, x_norm, y) }) + .collect(), + 16 => batch + .chunks_exact(16) + .map(|y| unsafe { cosine_once_x86::cosine_once_16_avx512(x, x_norm, y) }) + .collect(), + _ => batch + .chunks_exact(dimension) + .map(|y| unsafe { super::f32_x86::cosine_fast_avx512(x, x_norm, y) }) + .collect(), + } + } + + #[cfg(target_arch = "x86_64")] + #[target_feature(enable = "avx")] + pub(super) unsafe fn cosine_batch_avx( + x: &[f32], + x_norm: f32, + batch: &[f32], + dimension: usize, + ) -> Vec { + match dimension { + 8 => batch + .chunks_exact(8) + .map(|y| unsafe { cosine_once_x86::cosine_once_8_avx(x, x_norm, y) }) + .collect(), + 16 => batch + .chunks_exact(16) + .map(|y| unsafe { cosine_once_x86::cosine_once_16_avx(x, x_norm, y) }) + .collect(), + _ => batch + .chunks_exact(dimension) + .map(|y| unsafe { super::f32_x86::cosine_fast_avx(x, x_norm, y) }) + .collect(), + } + } +} + +/// Inlined f32 cosine kernels for builds whose baseline already guarantees AVX2 +/// (the default `haswell` wheel). No `#[target_feature]`, no runtime dispatch: +/// under `target-feature=+avx2,+fma` these compile to AVX2 and inline into the +/// batch loop exactly like the pre-PR code, so the modern path is not taxed by +/// the runtime-dispatch machinery (only needed below the AVX2 baseline). +#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] +mod f32_baseline { + use super::{dot, f32x8, f32x16, norm_l2}; + use crate::simd::{FloatSimd, SIMD}; + + #[inline] + pub fn cosine_once_8(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + unsafe { + let xv = f32x8::load_unaligned(x.as_ptr()); + let yv = f32x8::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + } + + #[inline] + pub fn cosine_once_16(x: &[f32], x_norm: f32, y: &[f32]) -> f32 { + unsafe { + let xv = f32x16::load_unaligned(x.as_ptr()); + let yv = f32x16::load_unaligned(y.as_ptr()); + let y2 = yv * yv; + let xy = xv * yv; + 1.0 - xy.reduce_sum() / x_norm / y2.reduce_sum().sqrt() + } + } + + #[inline] + pub fn cosine_fast(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { + unsafe { + let dim = x.len(); + let unrolled_len = dim / 16 * 16; + let mut y_norm16 = f32x16::zeros(); + let mut xy16 = f32x16::zeros(); + for i in (0..unrolled_len).step_by(16) { + let xv = f32x16::load_unaligned(x.as_ptr().add(i)); + let yv = f32x16::load_unaligned(other.as_ptr().add(i)); + xy16.multiply_add(xv, yv); + y_norm16.multiply_add(yv, yv); + } + let aligned_len = dim / 8 * 8; + let mut y_norm8 = f32x8::zeros(); + let mut xy8 = f32x8::zeros(); + for i in (unrolled_len..aligned_len).step_by(8) { + let xv = f32x8::load_unaligned(x.as_ptr().add(i)); + let yv = f32x8::load_unaligned(other.as_ptr().add(i)); + xy8.multiply_add(xv, yv); + y_norm8.multiply_add(yv, yv); + } + let y_norm = y_norm16.reduce_sum() + + y_norm8.reduce_sum() + + norm_l2(&other[aligned_len..]).powi(2); + let xy = xy16.reduce_sum() + + xy8.reduce_sum() + + dot(&x[aligned_len..], &other[aligned_len..]); + 1.0 - xy / x_norm / y_norm.sqrt() + } + } } impl Cosine for f32 { @@ -382,6 +524,7 @@ impl Cosine for f32 { cosine_with_norms_f32_dispatched(x, x_norm, y_norm, y) } + #[allow(unreachable_code)] fn cosine_batch<'a>( x: &'a [Self], batch: &'a [Self], @@ -389,6 +532,73 @@ impl Cosine for f32 { ) -> Box + 'a> { let x_norm = norm_l2(x); + // On a build whose baseline already guarantees AVX2 (the default + // `haswell` wheel), avoid the per-vector runtime dispatch + `#[target_feature]` + // wrapping that taxes the modern path. Dispatch ONCE per batch: AVX-512 + // hosts get the wide kernel; everyone else uses the inlined AVX2 baseline + // path (base-equivalent). The runtime-dispatch path below is only + // compiled/reached when the baseline is below AVX2 (pre-Haswell builds). + #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] + { + // dim 8/16 always use the inlined AVX2 baseline: AVX-512 gives no + // benefit for such tiny vectors (a masked 512-bit load is slower than + // a plain AVX2 load) and only adds dispatch + eager-collect overhead. + // Only the larger-dim path routes to AVX-512 on capable hosts — that's + // where the wider lanes actually pay off. + return match dimension { + 8 => Box::new( + batch + .chunks_exact(8) + .map(move |y| f32_baseline::cosine_once_8(x, x_norm, y)), + ), + 16 => Box::new( + batch + .chunks_exact(16) + .map(move |y| f32_baseline::cosine_once_16(x, x_norm, y)), + ), + _ => { + if matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { + Box::new( + unsafe { f32::cosine_batch_avx512(x, x_norm, batch, dimension) } + .into_iter(), + ) + } else { + Box::new( + batch + .chunks_exact(dimension) + .map(move |y| f32_baseline::cosine_fast(x, x_norm, y)), + ) + } + } + }; + } + + // Sub-AVX2 / non-x86 build: hoisted per-batch runtime dispatch. + #[cfg(target_arch = "x86_64")] + { + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => { + return Box::new( + unsafe { f32::cosine_batch_avx512(x, x_norm, batch, dimension) } + .into_iter(), + ); + } + SimdSupport::Avx2 | SimdSupport::AvxFma => { + return Box::new( + unsafe { f32::cosine_batch_avx_fma(x, x_norm, batch, dimension) } + .into_iter(), + ); + } + SimdSupport::Avx => { + return Box::new( + unsafe { f32::cosine_batch_avx(x, x_norm, batch, dimension) }.into_iter(), + ); + } + _ => {} + } + } + + // Scalar / non-x86 fallback. match dimension { 8 => Box::new( batch From d39482cf751dc0c2d40fb83cddd6a19cc63a8d61 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Wed, 8 Jul 2026 02:45:46 -0400 Subject: [PATCH 03/15] test(lance-linalg): cover batch cosine kernels and scalar x86 gather The runtime-dispatch codecov run flagged the f32 cosine batch kernels and the scalar x86 gather fallback as uncovered. Both are only reached at runtime on sub-AVX2 hosts, so on the AVX2+ CI runner they were never exercised even though the surrounding per-vector kernels have full parity tests. Add direct parity tests that call cosine_batch_avx, cosine_batch_avx_fma, and cosine_batch_avx512 across the 8/16/general dimension arms, comparing each batch entry against the scalar cosine_fast reference and gating each tier on the matching is_x86_feature_detected! check. Add a direct test for gather_scalar_x86 so the non-AVX2 gather path is covered on hosts that route through the AVX2 gather at runtime. --- rust/lance-linalg/src/distance/cosine.rs | 59 ++++++++++++++++++++++++ rust/lance-linalg/src/simd/f32.rs | 18 ++++++++ 2 files changed, 77 insertions(+) diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index c6615fa6757..bf2c5e6d6d6 100644 --- a/rust/lance-linalg/src/distance/cosine.rs +++ b/rust/lance-linalg/src/distance/cosine.rs @@ -1718,4 +1718,63 @@ mod tests { prop_assert!(approx::relative_eq!(scalar, avx, max_relative = 1e-5)); } } + + /// Asserts a batch-level f32 cosine SIMD kernel matches the scalar + /// `cosine_fast` reference for every vector in a multi-vector batch. Runs + /// each of the kernel's three internal dimension arms (8, 16, and the + /// general `chunks_exact` path). The batch kernels only run at runtime on + /// sub-AVX2 builds, so a direct call is the only way they get covered. + #[cfg(target_arch = "x86_64")] + fn check_cosine_batch_kernel(kernel: unsafe fn(&[f32], f32, &[f32], usize) -> Vec) { + for dimension in [8_usize, 16, 40] { + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.5 + 1.0).collect(); + let x_norm = norm_l2(&x); + let num_vectors = 3; + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 7) as f32) + 1.0) + .collect(); + + let got = unsafe { kernel(&x, x_norm, &batch, dimension) }; + assert_eq!(got.len(), num_vectors); + + let x_f64: Vec = x.iter().map(|&v| v as f64).collect(); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + let y_f64: Vec = chunk.iter().map(|&v| v as f64).collect(); + let expected = cosine_fast_scalar(&x_f64, x_norm, &y_f64); + assert_relative_eq!(g, expected, max_relative = 1e-3, epsilon = 1e-6); + } + } + } + + /// AVX + FMA batch kernel parity (AVX2 / AVX+FMA tiers). Runs on any + /// Haswell-or-newer host; early-returns without AVX and FMA. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_batch_avx_fma_matches_scalar() { + if !(std::is_x86_feature_detected!("avx") && std::is_x86_feature_detected!("fma")) { + return; + } + check_cosine_batch_kernel(super::f32::cosine_batch_avx_fma); + } + + /// AVX-only batch kernel parity (Sandy Bridge / Ivy Bridge tier). + /// Early-returns on hosts without AVX. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_batch_avx_matches_scalar() { + if !std::is_x86_feature_detected!("avx") { + return; + } + check_cosine_batch_kernel(super::f32::cosine_batch_avx); + } + + /// AVX-512 batch kernel parity. Early-returns on hosts without AVX-512F. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_cosine_batch_avx512_matches_scalar() { + if !std::is_x86_feature_detected!("avx512f") { + return; + } + check_cosine_batch_kernel(super::f32::cosine_batch_avx512); + } } diff --git a/rust/lance-linalg/src/simd/f32.rs b/rust/lance-linalg/src/simd/f32.rs index 638850af172..6ee3916751b 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -1084,4 +1084,22 @@ mod tests { let v = f32x8::gather(&a, &idx); assert_eq!(v.reduce_sum(), 113.0); } + + /// Directly exercises `gather_scalar_x86`, the per-index scalar fallback + /// `f32x8::gather` takes on x86_64 hosts without AVX2. Runtime AVX2 hosts + /// route through `gather_avx2` instead, so the fallback is otherwise never + /// hit under coverage. Reading the `__m256`-backed result needs AVX, so + /// skip on hosts without it. + #[cfg(target_arch = "x86_64")] + #[test] + fn test_gather_scalar_x86() { + if !std::is_x86_feature_detected!("avx") { + return; + } + let a = (0..256).map(|f| f as f32).collect::>(); + let idx = [0_i32, 4, 8, 12, 16, 20, 24, 29]; + let v = gather_scalar_x86(&a, &idx); + let expected = idx.map(|i| a[i as usize]); + assert_eq!(v.as_array(), expected); + } } From d000a67583fd78eb9d3eb594cb93b8fa71e99729 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:56:34 -0400 Subject: [PATCH 04/15] chore: address review feedback on runtime SIMD dispatch - cpu.rs: add a runnable doc example to simd_info() and link SimdInfo back to it - cpu.rs: pre-size the host-feature Vec, which takes up to 17 pushes - cpu.rs: parameterize the SimdSupport Display test with rstest cases - norm_l2.rs: hsum256_ps reduces via movehl/shuffle plus scalar adds, not two vhaddps; match the wording already used by the dot.rs and l2.rs siblings - simd/f32.rs: hoist as_array() out of the f32x16::find scan loop so the 16 lanes are materialized once instead of once per iteration - CONTRIBUTING.md: the workspace baseline is target-cpu=haswell, not x86-64-haswell --- CONTRIBUTING.md | 2 +- python/src/lib.rs | 8 ++++ rust/lance-core/src/utils/cpu.rs | 45 +++++++++++++++-------- rust/lance-linalg/src/distance/norm_l2.rs | 4 +- rust/lance-linalg/src/simd/f32.rs | 3 +- 5 files changed, 44 insertions(+), 18 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d96c7e451fb..3940ed7b683 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -27,7 +27,7 @@ Currently Lance is implemented in Rust and comes with a Python wrapper. So you'l ## Building for legacy x86_64 hosts (pre-Haswell) -The default workspace build targets `x86-64-haswell` (AVX2 + FMA + F16C), matching the published wheels. To build a binary that runs on pre-Haswell silicon (Sandy Bridge / Ivy Bridge / Westmere on Intel, Bulldozer / Piledriver / Steamroller on AMD — i.e. CPUs without AVX2), set the baseline yourself at build time: +The default workspace build targets `haswell` (AVX2 + FMA + F16C), matching the published wheels. To build a binary that runs on pre-Haswell silicon (Sandy Bridge / Ivy Bridge / Westmere on Intel, Bulldozer / Piledriver / Steamroller on AMD — i.e. CPUs without AVX2), set the baseline yourself at build time: ```sh RUSTFLAGS="-C target-cpu=x86-64-v2" cargo build --release diff --git a/python/src/lib.rs b/python/src/lib.rs index 208fa917e01..c5631d26f10 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -361,6 +361,14 @@ fn iops_counter() -> PyResult { /// "target_arch": str, # e.g. "x86_64", "aarch64", "loongarch64" /// "host_features": list[str], # raw CPU feature flags (x86_64 only) /// } +/// +/// Examples: +/// >>> import lance +/// >>> info = lance.simd_info() +/// >>> sorted(info) +/// ['host_features', 'target_arch', 'tier'] +/// >>> isinstance(info["tier"], str) +/// True #[pyfunction] pub fn simd_info(py: Python<'_>) -> PyResult> { let info = lance_core::utils::cpu::simd_info(); diff --git a/rust/lance-core/src/utils/cpu.rs b/rust/lance-core/src/utils/cpu.rs index d3b200f671f..5ac753d66e5 100644 --- a/rust/lance-core/src/utils/cpu.rs +++ b/rust/lance-core/src/utils/cpu.rs @@ -52,7 +52,7 @@ impl fmt::Display for SimdSupport { /// /// Mirrors the role of `pyarrow.runtime_info()`: a single, cheap call users can /// make to verify which SIMD tier the runtime selected and what underlying -/// features the host advertises. +/// features the host advertises. Obtain one with [`simd_info()`]. #[derive(Debug, Clone)] pub struct SimdInfo { /// The SIMD tier lance dispatches to at runtime on this host. @@ -69,7 +69,17 @@ pub struct SimdInfo { /// the raw CPU feature flags that drove the decision. /// /// Useful for performance debugging and giving users a way to verify which -/// dispatch tier they are hitting without rebuilding lance. +/// dispatch tier they are hitting without rebuilding lance. See [`SimdInfo`] +/// for the meaning of each field and [`SimdSupport`] for the tier values. +/// +/// # Examples +/// +/// ``` +/// use lance_core::utils::cpu::simd_info; +/// +/// let info = simd_info(); +/// println!("dispatching to {} on {}", info.tier, info.target_arch); +/// ``` pub fn simd_info() -> SimdInfo { SimdInfo { tier: *SIMD_SUPPORT, @@ -83,7 +93,7 @@ fn detect_host_features() -> Vec<&'static str> { // Each call must be inline: `is_x86_feature_detected!` does its own custom // input parsing and rejects feature names received via a `macro_rules!` // `:literal` metavariable on some toolchains. - let mut features = Vec::new(); + let mut features = Vec::with_capacity(17); if is_x86_feature_detected!("sse2") { features.push("sse2"); } @@ -275,6 +285,7 @@ mod aarch64 { #[cfg(test)] mod tests { use super::*; + use rstest::rstest; #[test] fn simd_info_exposes_tier() { @@ -300,17 +311,21 @@ mod tests { assert!(info.host_features.is_empty()); } - #[test] - fn simd_support_display_matches_lowercase_convention() { - assert_eq!(SimdSupport::Avx2.to_string(), "avx2"); - assert_eq!(SimdSupport::AvxFma.to_string(), "avx_fma"); - assert_eq!(SimdSupport::Avx.to_string(), "avx"); - assert_eq!(SimdSupport::Avx512.to_string(), "avx512"); - assert_eq!(SimdSupport::Avx512FP16.to_string(), "avx512_fp16"); - assert_eq!(SimdSupport::None.to_string(), "none"); - assert_eq!(SimdSupport::Neon.to_string(), "neon"); - assert_eq!(SimdSupport::Sse.to_string(), "sse"); - assert_eq!(SimdSupport::Lsx.to_string(), "lsx"); - assert_eq!(SimdSupport::Lasx.to_string(), "lasx"); + #[rstest] + #[case::none(SimdSupport::None, "none")] + #[case::neon(SimdSupport::Neon, "neon")] + #[case::sse(SimdSupport::Sse, "sse")] + #[case::avx(SimdSupport::Avx, "avx")] + #[case::avx_fma(SimdSupport::AvxFma, "avx_fma")] + #[case::avx2(SimdSupport::Avx2, "avx2")] + #[case::avx512(SimdSupport::Avx512, "avx512")] + #[case::avx512_fp16(SimdSupport::Avx512FP16, "avx512_fp16")] + #[case::lsx(SimdSupport::Lsx, "lsx")] + #[case::lasx(SimdSupport::Lasx, "lasx")] + fn simd_support_display_matches_lowercase_convention( + #[case] tier: SimdSupport, + #[case] expected: &str, + ) { + assert_eq!(tier.to_string(), expected); } } diff --git a/rust/lance-linalg/src/distance/norm_l2.rs b/rust/lance-linalg/src/distance/norm_l2.rs index fa89f2ce44d..25bbc121655 100644 --- a/rust/lance-linalg/src/distance/norm_l2.rs +++ b/rust/lance-linalg/src/distance/norm_l2.rs @@ -287,7 +287,9 @@ mod x86 { } /// Horizontal sum of an `__m256` register. Folds the upper 128-bit lane - /// into the lower, then sums lanes pairwise via two `vhaddps`. + /// into the lower, then sums lanes pairwise. Same shape as the helper in + /// the sibling `dot.rs` and `l2.rs` mod x86; kept local rather than + /// hoisted to a shared module to avoid a one-helper module file. #[inline] #[target_feature(enable = "avx")] unsafe fn hsum256_ps(v: __m256) -> f32 { diff --git a/rust/lance-linalg/src/simd/f32.rs b/rust/lance-linalg/src/simd/f32.rs index 6ee3916751b..974347443e7 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -735,8 +735,9 @@ impl SIMD for f32x16 { unsafe { // _mm256_cmpeq_ps_mask requires AVX-512 (avx512f); use a scalar scan here // since we only require AVX2. + let arr = self.as_array(); for i in 0..16 { - if self.as_array().get_unchecked(i) == &val { + if arr.get_unchecked(i) == &val { return Some(i as i32); } } From d5205dc9a197fc501893fa324a32a49a4713c887 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:07:22 -0400 Subject: [PATCH 05/15] refactor: share the AVX horizontal-sum helper across the distance kernels - simd/x86.rs: new crate-internal module holding hsum256_ps and hsum256_pd. The f32 reduction had been copied verbatim into cosine.rs, dot.rs, l2.rs and norm_l2.rs, and each copy carried a doc comment explaining why the *other* copies existed. One of them had already drifted from its implementation. Bodies are unchanged; the helpers now have direct unit tests. - simd/f32.rs, simd/f64.rs: the test feature guards checked avx2 while their comments said AVX. Neither file reaches an AVX2 intrinsic outside the runtime-detected gather, so the guards now check what the tests actually need: avx, plus fma for the ones that call multiply_add (which lowers to _mm256_fmadd_ps). Guarding those on avx alone would have executed vfmadd on a Sandy Bridge host. This matches the guards the distance kernels already use, and lets AVX-only and AVX+FMA hosts exercise the tests. - l2.rs: the SimdSupport import was unconditional, but no variant is named unless fp16kernels is on or the target is x86_64, so `cargo clippy --all --tests --benches -- -D warnings` failed on aarch64. Gate the import on the condition instead of silencing the lint. --- rust/lance-linalg/src/distance/cosine.rs | 41 ++-------- rust/lance-linalg/src/distance/dot.rs | 16 +--- rust/lance-linalg/src/distance/l2.rs | 20 +---- rust/lance-linalg/src/distance/norm_l2.rs | 16 +--- rust/lance-linalg/src/simd.rs | 2 + rust/lance-linalg/src/simd/f32.rs | 24 +++--- rust/lance-linalg/src/simd/f64.rs | 21 +++-- rust/lance-linalg/src/simd/x86.rs | 96 +++++++++++++++++++++++ 8 files changed, 138 insertions(+), 98 deletions(-) create mode 100644 rust/lance-linalg/src/simd/x86.rs diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index bf2c5e6d6d6..457c0c0dd6a 100644 --- a/rust/lance-linalg/src/distance/cosine.rs +++ b/rust/lance-linalg/src/distance/cosine.rs @@ -664,6 +664,7 @@ mod f64_x86 { use std::arch::x86_64::*; use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_pd; use crate::simd::{FloatSimd, SIMD}; /// AVX-512 path for f64 fast cosine: 8-wide `__m512d` xy/yy with `vfmadd231pd` per iteration. @@ -742,8 +743,8 @@ mod f64_x86 { acc_yy = _mm256_add_pd(acc_yy, _mm256_mul_pd(yv, yv)); } - let xy_main = hsum256_pd_local(acc_xy); - let yy_main = hsum256_pd_local(acc_yy); + let xy_main = hsum256_pd(acc_xy); + let yy_main = hsum256_pd(acc_yy); let tail_y_norm: f64 = y[aligned_len..].iter().map(|&v| v * v).sum(); let tail_xy: f64 = x[aligned_len..] @@ -756,19 +757,6 @@ mod f64_x86 { let xy = (xy_main + tail_xy) as f32; 1.0 - xy / x_norm / y_norm_sq.sqrt() } - - /// Horizontal sum of an `__m256d` register. Local to this `f64_x86` - /// mod for cache locality with its inner kernels; matches the - /// per-mod hsum-helper pattern already used in `dot.rs`/`norm_l2.rs`. - #[inline] - #[target_feature(enable = "avx")] - unsafe fn hsum256_pd_local(v: __m256d) -> f64 { - let lo = _mm256_castpd256_pd128(v); - let hi = _mm256_extractf128_pd(v, 1); - let sum128 = _mm_add_pd(lo, hi); - let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); - _mm_cvtsd_f64(sum64) - } } #[cfg(not(target_arch = "x86_64"))] @@ -897,6 +885,7 @@ mod f32_x86 { use std::arch::x86_64::*; use super::{dot, f32x8, f32x16, norm_l2}; + use crate::simd::x86::hsum256_ps; use crate::simd::{FloatSimd, SIMD}; /// AVX + FMA path for f32 fast cosine. Covers both AvxFma and AVX2 dispatch (body uses no AVX2-specific intrinsics). @@ -943,30 +932,14 @@ mod f32_x86 { acc_yy = _mm256_add_ps(acc_yy, _mm256_mul_ps(yv, yv)); } - let xy_main = hsum256_ps_local(acc_xy); - let yy_main = hsum256_ps_local(acc_yy); + let xy_main = hsum256_ps(acc_xy); + let yy_main = hsum256_ps(acc_yy); let y_norm = yy_main + norm_l2(&other[aligned_len..]).powi(2); let xy = xy_main + dot(&x[aligned_len..], &other[aligned_len..]); 1.0 - xy / x_norm / y_norm.sqrt() } - /// Horizontal sum of an `__m256` register. Local to this `f32_x86` mod - /// because the sibling `mod x86` in `norm_l2.rs` keeps its own copy - /// for cache locality with its inner kernels; duplicating the helper - /// here avoids a cross-module unsafe re-export and matches the pattern - /// already used in `dot.rs::x86`. - #[inline] - #[target_feature(enable = "avx")] - unsafe fn hsum256_ps_local(v: __m256) -> f32 { - let lo = _mm256_castps256_ps128(v); - let hi = _mm256_extractf128_ps(v, 1); - let sum128 = _mm_add_ps(lo, hi); - let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); - let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); - _mm_cvtss_f32(sum32) - } - /// AVX-512 path for f32 fast cosine: 16-wide `__m512` xy/yy with `vfmadd231ps` per iteration. #[target_feature(enable = "avx512f")] pub unsafe fn cosine_fast_avx512(x: &[f32], x_norm: f32, other: &[f32]) -> f32 { @@ -1048,7 +1021,7 @@ mod f32_x86 { acc = _mm256_add_ps(acc, _mm256_mul_ps(xv, yv)); } - let xy_main = hsum256_ps_local(acc); + let xy_main = hsum256_ps(acc); let xy = xy_main + dot(&x[aligned_len..], &y[aligned_len..]); 1.0 - xy / x_norm / y_norm } diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 7d3e7a36ff0..d39ee2f0f48 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -304,6 +304,7 @@ mod x86 { use std::arch::x86_64::*; use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_ps; use crate::simd::{FloatSimd, SIMD}; /// AVX-512 path for f64: 8-wide `__m512d` with `vfmadd231pd` per iteration. @@ -387,21 +388,6 @@ mod x86 { (acc_sum + tail) as f32 } - /// Horizontal sum of an `__m256` register. Folds the upper 128-bit lane - /// into the lower, then sums lanes pairwise. Same shape as the helper in - /// the sibling `norm_l2.rs` mod x86; kept local rather than hoisted to a - /// shared module to avoid a one-helper module file. - #[inline] - #[target_feature(enable = "avx")] - unsafe fn hsum256_ps(v: __m256) -> f32 { - let lo = _mm256_castps256_ps128(v); - let hi = _mm256_extractf128_ps(v, 1); - let sum128 = _mm_add_ps(lo, hi); - let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); - let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); - _mm_cvtss_f32(sum32) - } - /// AVX-512 path for f32: 16-wide `__m512` with `vfmadd231ps` per iteration. #[target_feature(enable = "avx512f")] pub unsafe fn dot_f32_avx512(x: &[f32], y: &[f32]) -> f32 { diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 78886860401..40687e33f85 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -20,6 +20,9 @@ use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray}; use lance_core::assume_eq; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::cpu::SIMD_SUPPORT; +// Named tiers are only matched on x86_64, or by the fp16 kernels on the other +// architectures; without either, nothing below names a `SimdSupport` variant. +#[cfg(any(feature = "fp16kernels", target_arch = "x86_64"))] use lance_core::utils::cpu::SimdSupport; use num_traits::{AsPrimitive, Num}; @@ -51,7 +54,6 @@ pub fn l2(from: &[T], to: &[T]) -> f32 { pub fn l2_f32(x: &[f32], y: &[f32]) -> f32 { #[cfg(target_arch = "x86_64")] { - use lance_core::utils::cpu::SimdSupport; if matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { // SAFETY: guarded by the runtime AVX-512 detection above. return unsafe { l2_f32_avx512(x, y) }; @@ -341,6 +343,7 @@ mod x86 { use std::arch::x86_64::*; use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_ps; use crate::simd::{FloatSimd, SIMD}; /// AVX-512 path for f64: 8-wide `__m512d` with `vsubpd` + `vfmadd231pd` per iteration. @@ -437,21 +440,6 @@ mod x86 { (acc_sum + tail) as f32 } - /// Horizontal sum of an `__m256` register. Folds the upper 128-bit lane - /// into the lower, then sums lanes pairwise. Same shape as the helper in - /// the sibling `norm_l2.rs` mod x86; kept local rather than hoisted to a - /// shared module to avoid a one-helper module file. - #[inline] - #[target_feature(enable = "avx")] - unsafe fn hsum256_ps(v: __m256) -> f32 { - let lo = _mm256_castps256_ps128(v); - let hi = _mm256_extractf128_ps(v, 1); - let sum128 = _mm_add_ps(lo, hi); - let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); - let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); - _mm_cvtss_f32(sum32) - } - /// AVX-512 path for f32: 16-wide `__m512` with `vsubps` + `vfmadd231ps` per iteration. #[target_feature(enable = "avx512f")] pub unsafe fn l2_f32_avx512(x: &[f32], y: &[f32]) -> f32 { diff --git a/rust/lance-linalg/src/distance/norm_l2.rs b/rust/lance-linalg/src/distance/norm_l2.rs index 25bbc121655..8d126d6ed57 100644 --- a/rust/lance-linalg/src/distance/norm_l2.rs +++ b/rust/lance-linalg/src/distance/norm_l2.rs @@ -222,6 +222,7 @@ mod x86 { use std::arch::x86_64::*; use crate::simd::f64::{f64x4, f64x8}; + use crate::simd::x86::hsum256_ps; use crate::simd::{FloatSimd, SIMD}; /// AVX-512 path for f64: 8-wide `__m512d` with `vfmadd231pd` per iteration. @@ -286,21 +287,6 @@ mod x86 { (acc_sum + tail).sqrt() as f32 } - /// Horizontal sum of an `__m256` register. Folds the upper 128-bit lane - /// into the lower, then sums lanes pairwise. Same shape as the helper in - /// the sibling `dot.rs` and `l2.rs` mod x86; kept local rather than - /// hoisted to a shared module to avoid a one-helper module file. - #[inline] - #[target_feature(enable = "avx")] - unsafe fn hsum256_ps(v: __m256) -> f32 { - let lo = _mm256_castps256_ps128(v); - let hi = _mm256_extractf128_ps(v, 1); - let sum128 = _mm_add_ps(lo, hi); - let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); - let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); - _mm_cvtss_f32(sum32) - } - /// AVX-512 path for f32: 16-wide `__m512` with `vfmadd231ps` per iteration. #[target_feature(enable = "avx512f")] pub unsafe fn norm_l2_f32_avx512(vector: &[f32]) -> f32 { diff --git a/rust/lance-linalg/src/simd.rs b/rust/lance-linalg/src/simd.rs index 91dc1c6959d..6722eafd768 100644 --- a/rust/lance-linalg/src/simd.rs +++ b/rust/lance-linalg/src/simd.rs @@ -19,6 +19,8 @@ pub mod f32; pub mod f64; pub mod i32; pub mod u8; +#[cfg(target_arch = "x86_64")] +pub(crate) mod x86; use num_traits::{Float, Num}; use u8::u8x16; diff --git a/rust/lance-linalg/src/simd/f32.rs b/rust/lance-linalg/src/simd/f32.rs index 974347443e7..11a7f7baa37 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -927,10 +927,11 @@ mod tests { #[test] fn test_basic_ops() { - // The `f32x8` constructor / load / store paths are AVX-only on x86_64 - // after the baseline lowering; skip on hosts that don't support AVX2. + // Load / store / arithmetic on `f32x8` lower to AVX intrinsics, and + // `multiply_add` lowers to `_mm256_fmadd_ps`, which needs FMA. Both + // are present from the AvxFma tier up. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { return; } let a = (0..8).map(|f| f as f32).collect::>(); @@ -970,8 +971,9 @@ mod tests { #[test] fn test_f32x8_cmp_ops() { + // `min` / `reduce_min` are AVX intrinsics; `find` is a scalar scan. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") { return; } let a = [1.0_f32, 2.0, 5.0, 6.0, 7.0, 3.0, 2.0, 1.0]; @@ -998,8 +1000,9 @@ mod tests { #[test] fn test_basic_f32x16_ops() { + // `f32x16` is a pair of `__m256`; `multiply_add` needs FMA. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { return; } let a = (0..16).map(|f| f as f32).collect::>(); @@ -1036,8 +1039,9 @@ mod tests { #[test] fn test_f32x16_cmp_ops() { + // `min` / `reduce_min` are AVX intrinsics; `find` is a scalar scan. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") { return; } let a = [ @@ -1073,11 +1077,11 @@ mod tests { #[test] fn test_f32x8_gather() { - // `f32x8::gather` does its own runtime AVX2 detection, but the - // returned `f32x8` itself is `__m256`-backed on x86_64 and the - // `reduce_sum` helper requires AVX. Skip on pre-Haswell hosts. + // `f32x8::gather` does its own runtime AVX2 detection and falls back + // to a scalar gather, so this test only needs whatever reading the + // `__m256`-backed result costs: AVX, for `reduce_sum`. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") { return; } let a = (0..256).map(|f| f as f32).collect::>(); diff --git a/rust/lance-linalg/src/simd/f64.rs b/rust/lance-linalg/src/simd/f64.rs index 0274c8ad6f1..1276f54da56 100644 --- a/rust/lance-linalg/src/simd/f64.rs +++ b/rust/lance-linalg/src/simd/f64.rs @@ -736,10 +736,10 @@ mod tests { #[test] fn test_f64x4_basic_ops() { - // The `f64x4` constructor / load / store paths are AVX-only on x86_64 - // after the baseline lowering; skip on hosts that don't support AVX2. + // The `f64x4` constructor / load / store / arithmetic paths all lower + // to AVX intrinsics on x86_64; none of them need AVX2. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") { return; } let a = [1.0_f64, 2.0, 3.0, 4.0]; @@ -763,8 +763,9 @@ mod tests { #[test] fn test_f64x4_fma() { + // `multiply_add` lowers to `_mm256_fmadd_pd`, which needs FMA. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { return; } let a = [1.0_f64, 2.0, 3.0, 4.0]; @@ -779,8 +780,9 @@ mod tests { #[test] fn test_f64x4_min() { + // `min` / `reduce_min` are AVX intrinsics. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") { return; } let a = [1.0_f64, 5.0, 2.0, 8.0]; @@ -795,8 +797,9 @@ mod tests { #[test] fn test_f64x8_basic_ops() { + // `f64x8` is a pair of `__m256d`; add / reduce are AVX intrinsics. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") { return; } let a: [f64; 8] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; @@ -817,8 +820,9 @@ mod tests { #[test] fn test_f64x8_fma() { + // `multiply_add` lowers to `_mm256_fmadd_pd`, which needs FMA. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { return; } let a: [f64; 8] = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; @@ -834,8 +838,9 @@ mod tests { #[test] fn test_f64x8_min() { + // `min` / `reduce_min` are AVX intrinsics. #[cfg(target_arch = "x86_64")] - if !std::is_x86_feature_detected!("avx2") { + if !std::is_x86_feature_detected!("avx") { return; } let a: [f64; 8] = [5.0, 1.0, 8.0, 3.0, 9.0, 2.0, 7.0, 4.0]; diff --git a/rust/lance-linalg/src/simd/x86.rs b/rust/lance-linalg/src/simd/x86.rs new file mode 100644 index 00000000000..7abebd8b487 --- /dev/null +++ b/rust/lance-linalg/src/simd/x86.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Reduction helpers shared by the x86_64 AVX kernels in [`crate::distance`]. +//! +//! Each distance kernel accumulates into a 256-bit register and folds it down +//! to a scalar once, after the main loop. The fold is identical across +//! `cosine`, `dot`, `l2` and `norm_l2`, so it lives here instead of being +//! copied into each kernel's private `mod x86`. + +use std::arch::x86_64::*; + +/// Horizontal sum of the eight `f32` lanes of an `__m256`. +/// +/// Folds the upper 128-bit lane into the lower one, then reduces the +/// remaining four lanes pairwise. Uses `movehl`/`shuffle` plus scalar adds +/// rather than two `vhaddps`, which is one fewer uop on most cores. +/// +/// # Safety +/// +/// The host must support AVX. Callers are `#[target_feature]`-annotated +/// kernels that the runtime dispatcher only selects after checking. +#[inline] +#[target_feature(enable = "avx")] +pub unsafe fn hsum256_ps(v: __m256) -> f32 { + let lo = _mm256_castps256_ps128(v); + let hi = _mm256_extractf128_ps(v, 1); + let sum128 = _mm_add_ps(lo, hi); + let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); + _mm_cvtss_f32(sum32) +} + +/// Horizontal sum of the four `f64` lanes of an `__m256d`. +/// +/// Folds the upper 128-bit lane into the lower one, then adds the remaining +/// pair. +/// +/// # Safety +/// +/// The host must support AVX. Callers are `#[target_feature]`-annotated +/// kernels that the runtime dispatcher only selects after checking. +#[inline] +#[target_feature(enable = "avx")] +pub unsafe fn hsum256_pd(v: __m256d) -> f64 { + let lo = _mm256_castpd256_pd128(v); + let hi = _mm256_extractf128_pd(v, 1); + let sum128 = _mm_add_pd(lo, hi); + let sum64 = _mm_add_pd(sum128, _mm_unpackhi_pd(sum128, sum128)); + _mm_cvtsd_f64(sum64) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hsum256_ps_sums_all_eight_lanes() { + if !std::is_x86_feature_detected!("avx") { + return; + } + let v = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let sum = unsafe { hsum256_ps(_mm256_loadu_ps(v.as_ptr())) }; + assert_eq!(sum, 36.0); + } + + #[test] + fn hsum256_ps_handles_negative_lanes() { + if !std::is_x86_feature_detected!("avx") { + return; + } + let v = [-1.5_f32, 2.0, -3.0, 4.5, 0.0, -0.5, 1.0, 2.5]; + let sum = unsafe { hsum256_ps(_mm256_loadu_ps(v.as_ptr())) }; + assert_eq!(sum, 5.0); + } + + #[test] + fn hsum256_pd_sums_all_four_lanes() { + if !std::is_x86_feature_detected!("avx") { + return; + } + let v = [1.0_f64, 2.0, 3.0, 4.0]; + let sum = unsafe { hsum256_pd(_mm256_loadu_pd(v.as_ptr())) }; + assert_eq!(sum, 10.0); + } + + #[test] + fn hsum256_pd_handles_negative_lanes() { + if !std::is_x86_feature_detected!("avx") { + return; + } + let v = [-1.5_f64, 2.0, -3.0, 4.5]; + let sum = unsafe { hsum256_pd(_mm256_loadu_pd(v.as_ptr())) }; + assert_eq!(sum, 2.0); + } +} From bdeda9783135bbd170d010f5c1640a73c305b010 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:02:06 -0400 Subject: [PATCH 06/15] test: parameterize the AVX horizontal-sum tests with rstest The four hand-rolled tests differed only in their input lanes and expected sum, which AGENTS.md asks to express as rstest cases. Fold them into two parameterized tests and add zero and cancelling-lane cases while here. Also record why the helpers are `pub` inside a `pub(crate)` module, and what the 0x55 shuffle mask selects. --- Cargo.lock | 1 + rust/lance-linalg/Cargo.toml | 1 + rust/lance-linalg/src/simd/x86.rs | 53 ++++++++++++++----------------- 3 files changed, 25 insertions(+), 30 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c72f4ed4410..b4efa176e3a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4932,6 +4932,7 @@ dependencies = [ "proptest", "rand 0.9.4", "rayon", + "rstest", ] [[package]] diff --git a/rust/lance-linalg/Cargo.toml b/rust/lance-linalg/Cargo.toml index 6a188ec3c62..d94efa7b2fa 100644 --- a/rust/lance-linalg/Cargo.toml +++ b/rust/lance-linalg/Cargo.toml @@ -25,6 +25,7 @@ approx = { workspace = true } criterion = { workspace = true } lance-testing = { path = "../lance-testing" } proptest.workspace = true +rstest.workspace = true [build-dependencies] cc = "1.0.83" diff --git a/rust/lance-linalg/src/simd/x86.rs b/rust/lance-linalg/src/simd/x86.rs index 7abebd8b487..437018669e2 100644 --- a/rust/lance-linalg/src/simd/x86.rs +++ b/rust/lance-linalg/src/simd/x86.rs @@ -7,6 +7,10 @@ //! to a scalar once, after the main loop. The fold is identical across //! `cosine`, `dot`, `l2` and `norm_l2`, so it lives here instead of being //! copied into each kernel's private `mod x86`. +//! +//! The module itself is `pub(crate)`, which is what keeps these helpers off the +//! public API; the items are `pub` rather than `pub(crate)` only because +//! `clippy::redundant_pub_crate` fires on the narrower visibility. use std::arch::x86_64::*; @@ -27,6 +31,8 @@ pub unsafe fn hsum256_ps(v: __m256) -> f32 { let hi = _mm256_extractf128_ps(v, 1); let sum128 = _mm_add_ps(lo, hi); let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128)); + // 0x55 broadcasts lane 1 into lane 0, so the scalar add below lands the + // last of the four partial sums. let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55)); _mm_cvtss_f32(sum32) } @@ -53,44 +59,31 @@ pub unsafe fn hsum256_pd(v: __m256d) -> f64 { #[cfg(test)] mod tests { use super::*; + use rstest::rstest; - #[test] - fn hsum256_ps_sums_all_eight_lanes() { - if !std::is_x86_feature_detected!("avx") { - return; - } - let v = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; - let sum = unsafe { hsum256_ps(_mm256_loadu_ps(v.as_ptr())) }; - assert_eq!(sum, 36.0); - } - - #[test] - fn hsum256_ps_handles_negative_lanes() { - if !std::is_x86_feature_detected!("avx") { - return; - } - let v = [-1.5_f32, 2.0, -3.0, 4.5, 0.0, -0.5, 1.0, 2.5]; - let sum = unsafe { hsum256_ps(_mm256_loadu_ps(v.as_ptr())) }; - assert_eq!(sum, 5.0); - } - - #[test] - fn hsum256_pd_sums_all_four_lanes() { + #[rstest] + #[case::ascending([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], 36.0)] + #[case::negative_lanes([-1.5, 2.0, -3.0, 4.5, 0.0, -0.5, 1.0, 2.5], 5.0)] + #[case::zeros([0.0; 8], 0.0)] + #[case::cancelling([1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0, -4.0], 0.0)] + fn hsum256_ps_sums_every_lane(#[case] lanes: [f32; 8], #[case] expected: f32) { if !std::is_x86_feature_detected!("avx") { return; } - let v = [1.0_f64, 2.0, 3.0, 4.0]; - let sum = unsafe { hsum256_pd(_mm256_loadu_pd(v.as_ptr())) }; - assert_eq!(sum, 10.0); + let sum = unsafe { hsum256_ps(_mm256_loadu_ps(lanes.as_ptr())) }; + assert_eq!(sum, expected); } - #[test] - fn hsum256_pd_handles_negative_lanes() { + #[rstest] + #[case::ascending([1.0, 2.0, 3.0, 4.0], 10.0)] + #[case::negative_lanes([-1.5, 2.0, -3.0, 4.5], 2.0)] + #[case::zeros([0.0; 4], 0.0)] + #[case::cancelling([1.0, -1.0, 2.0, -2.0], 0.0)] + fn hsum256_pd_sums_every_lane(#[case] lanes: [f64; 4], #[case] expected: f64) { if !std::is_x86_feature_detected!("avx") { return; } - let v = [-1.5_f64, 2.0, -3.0, 4.5]; - let sum = unsafe { hsum256_pd(_mm256_loadu_pd(v.as_ptr())) }; - assert_eq!(sum, 2.0); + let sum = unsafe { hsum256_pd(_mm256_loadu_pd(lanes.as_ptr())) }; + assert_eq!(sum, expected); } } From 1903ba47125e9cf9b9f81251f5bbb88065b28eaa Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:17:44 -0400 Subject: [PATCH 07/15] fix: bounds-check the scalar gather fallback on x86_64 gather_scalar_x86 offset a raw pointer by each index, so an out-of-range index read out of bounds instead of panicking. It is the path taken on hosts without AVX2, where the bounds checks cost nothing worth measuring. Index the slice instead, and cover the out-of-range case. --- rust/lance-linalg/src/simd/f32.rs | 33 ++++++++++++++++++------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/rust/lance-linalg/src/simd/f32.rs b/rust/lance-linalg/src/simd/f32.rs index 11a7f7baa37..f29d30665a7 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -118,23 +118,16 @@ unsafe fn gather_avx2(slice: &[f32], indices: &[i32; 8]) -> f32x8 { } /// Portable scalar gather for x86_64 hosts without AVX2. +/// +/// Indexes the slice rather than offsetting a raw pointer: this is the slow +/// path already, so an out-of-range index should panic instead of reading out +/// of bounds. #[cfg(target_arch = "x86_64")] #[inline] fn gather_scalar_x86(slice: &[f32], indices: &[i32; 8]) -> f32x8 { - let ptr = slice.as_ptr(); - unsafe { - let values = [ - *ptr.add(indices[0] as usize), - *ptr.add(indices[1] as usize), - *ptr.add(indices[2] as usize), - *ptr.add(indices[3] as usize), - *ptr.add(indices[4] as usize), - *ptr.add(indices[5] as usize), - *ptr.add(indices[6] as usize), - *ptr.add(indices[7] as usize), - ]; - f32x8::load_unaligned(values.as_ptr()) - } + let values = indices.map(|i| slice[i as usize]); + // SAFETY: `values` is eight contiguous, initialized `f32`. + unsafe { f32x8::load_unaligned(values.as_ptr()) } } impl From<&[f32]> for f32x8 { @@ -1107,4 +1100,16 @@ mod tests { let expected = idx.map(|i| a[i as usize]); assert_eq!(v.as_array(), expected); } + + /// An index past the end of the slice panics rather than reading out of + /// bounds. The bounds check fires before any AVX instruction, so this + /// case runs on every x86_64 host. + #[cfg(target_arch = "x86_64")] + #[test] + #[should_panic(expected = "index out of bounds")] + fn test_gather_scalar_x86_rejects_out_of_range_index() { + let a = (0..8).map(|f| f as f32).collect::>(); + let idx = [0_i32, 1, 2, 3, 4, 5, 6, 99]; + let _ = gather_scalar_x86(&a, &idx); + } } From c247ddd11bb0e2196f364dd28a0267cbfe7e4489 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:37:23 -0400 Subject: [PATCH 08/15] fix: validate gather indices once, before backend dispatch Bounds-checking only the scalar fallback left f32x8::gather with two behaviours: a panic on hosts without AVX2, and an out-of-bounds read on hosts with it, since vgatherdps checks nothing. The NEON and LASX arms offset a raw pointer and had the same problem. Check the indices once in gather() so every backend rejects a bad index the same way. It has no callers on a hot path, and eight comparisons are nothing next to the gather itself. --- rust/lance-linalg/src/simd/f32.rs | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/rust/lance-linalg/src/simd/f32.rs b/rust/lance-linalg/src/simd/f32.rs index f29d30665a7..4ce7f64706d 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -57,8 +57,24 @@ impl f32x8 { /// `multiversion` crate) was considered but doesn't fit here: the function /// returns `Self` and `_mm256_i32gather_ps::<4>` requires the const-generic /// stride to be a compile-time literal — neither composes with the macro. + /// + /// # Panics + /// + /// If any index is negative or lands outside `slice`. #[inline] pub fn gather(slice: &[f32], indices: &[i32; 8]) -> Self { + // Every backend below reads without bounds checking: `vgatherdps` does + // none, and the NEON / LASX arms offset a raw pointer. Check once here + // so an out-of-range index panics on every host rather than reading out + // of bounds on some and panicking on others. + for &i in indices { + assert!( + (i as usize) < slice.len(), + "gather index {i} is out of bounds for a slice of length {}", + slice.len() + ); + } + #[cfg(target_arch = "x86_64")] { if is_x86_feature_detected!("avx2") { @@ -917,6 +933,7 @@ impl SubAssign for f32x16 { mod tests { use super::*; + use rstest::rstest; #[test] fn test_basic_ops() { @@ -1112,4 +1129,17 @@ mod tests { let idx = [0_i32, 1, 2, 3, 4, 5, 6, 99]; let _ = gather_scalar_x86(&a, &idx); } + + /// `gather` validates before dispatching, so every backend — `vgatherdps`, + /// the x86 scalar fallback, and the NEON / LASX raw-pointer arms — rejects + /// a bad index identically instead of reading out of bounds. + #[rstest] + #[case::past_end(99)] + #[case::negative(-1)] + #[should_panic(expected = "out of bounds")] + fn test_gather_rejects_invalid_index(#[case] bad_index: i32) { + let a = (0..8).map(|f| f as f32).collect::>(); + let idx = [0_i32, 1, 2, 3, 4, 5, 6, bad_index]; + let _ = f32x8::gather(&a, &idx); + } } From bb8af46f5a8177f1ab8abfa17dd218dbc8cbf28e Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:26:41 -0400 Subject: [PATCH 09/15] fix: require FMA before selecting the AVX2 dispatch tier Every kernel the Avx2 tier dispatches to is annotated `#[target_feature(enable = "avx,fma")]`, but the tier was selected on `is_x86_feature_detected!("avx2")` alone. AVX2 does not imply FMA in the ISA. No shipping AVX2 part lacks FMA, so this was latent rather than live, but it is the same class of fault the rest of this work exists to prevent. Check FMA explicitly, document the tier's contract, and add a test that fails if a tier is ever selected on a host that cannot run its kernels. --- rust/lance-core/src/utils/cpu.rs | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/rust/lance-core/src/utils/cpu.rs b/rust/lance-core/src/utils/cpu.rs index 5ac753d66e5..c4d5a976cbb 100644 --- a/rust/lance-core/src/utils/cpu.rs +++ b/rust/lance-core/src/utils/cpu.rs @@ -20,6 +20,10 @@ pub enum SimdSupport { /// AVX + FMA but no AVX2. /// AMD Piledriver / Steamroller / FX-7500. AvxFma, + /// AVX2 + FMA. Intel Haswell / AMD Excavator and later. + /// + /// Selecting this tier asserts FMA is present: the kernels it dispatches to + /// are `#[target_feature(enable = "avx,fma")]`. Avx2, Avx512, Avx512FP16, @@ -179,7 +183,12 @@ pub static SIMD_SUPPORT: LazyLock = LazyLock::new(|| { } else { SimdSupport::Avx512 } - } else if is_x86_feature_detected!("avx2") { + } else if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { + // FMA is checked explicitly: every kernel selected for this tier is + // `#[target_feature(enable = "avx,fma")]`, and AVX2 does not imply + // FMA in the ISA. Every shipping AVX2 part has FMA, so this only + // guards against a host that would otherwise take an FMA kernel + // without FMA. SimdSupport::Avx2 } else if is_x86_feature_detected!("avx") && is_x86_feature_detected!("fma") { // AMD Piledriver / Steamroller / FX-7500: 256-bit float ops + FMA but no AVX2. @@ -311,6 +320,23 @@ mod tests { assert!(info.host_features.is_empty()); } + /// The `Avx2` and `AvxFma` tiers both dispatch to kernels declared + /// `#[target_feature(enable = "avx,fma")]`, so neither may be selected on a + /// host without FMA. AVX2 does not imply FMA in the ISA, so the detection + /// checks it explicitly. (`Avx512*` is excluded: its kernels declare + /// `avx512f`, which is what `has_avx512` verifies.) + #[cfg(target_arch = "x86_64")] + #[test] + fn avx_fma_tiers_are_only_selected_when_fma_is_detected() { + if matches!(*SIMD_SUPPORT, SimdSupport::Avx2 | SimdSupport::AvxFma) { + assert!( + is_x86_feature_detected!("fma"), + "tier {} dispatches to avx,fma kernels but the host has no FMA", + *SIMD_SUPPORT + ); + } + } + #[rstest] #[case::none(SimdSupport::None, "none")] #[case::neon(SimdSupport::Neon, "neon")] From 2c16954a14e50e32fbb0f408600f313a4575afc5 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:26:41 -0400 Subject: [PATCH 10/15] perf: choose the SIMD tier once per batch for f32 dot and l2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runtime dispatch was happening once per vector inside the batch loops. On a build whose baseline already implies AVX2 — the default `haswell` wheel — the old code inlined the kernel directly, so the `match *SIMD_SUPPORT` plus the non-inlinable `#[target_feature]` call bought nothing and cost a measured +24% at dim 1024 and +36% at dim 8 on Broadwell. Cosine was fixed this way already; dot and l2 were not. Dispatch once, for the whole batch: - On AVX2-baseline builds, call the same inlined auto-vectorized kernel the code used before dispatch existed. Only an AVX-512 host with dimension > 16 takes a wide kernel, entered once. - On sub-AVX2 builds, still dispatch, but into a kernel that loops internally. `Dot` had no batch hook at all, so `dot_distance_batch` mapped `dot_distance` per vector; add `Dot::dot_batch` with a per-vector default, so every other impl keeps its current behaviour. Both batch methods return `impl Iterator` rather than a trait object: the k-means assignment loop drives them one element at a time, and a `Box` would trade the per-vector dispatch for a per-element virtual call. `BatchIter` keeps the lazy and eager shapes in one statically dispatched type. `do_dot_distance_arrow_batch` carried its own copy of the per-vector loop and would not have benefited; route it through `dot_distance_batch`. Public signatures are unchanged from before the dispatch work: `l2_distance_batch` still returns `impl Iterator`, `dot_distance_batch` still returns a boxed one. --- rust/lance-linalg/src/distance.rs | 40 ++++ rust/lance-linalg/src/distance/dot.rs | 251 +++++++++++++++++++++++++- rust/lance-linalg/src/distance/l2.rs | 213 +++++++++++++++++++++- 3 files changed, 497 insertions(+), 7 deletions(-) diff --git a/rust/lance-linalg/src/distance.rs b/rust/lance-linalg/src/distance.rs index 23d1cae2d63..0a2cb31a7d4 100644 --- a/rust/lance-linalg/src/distance.rs +++ b/rust/lance-linalg/src/distance.rs @@ -25,6 +25,46 @@ pub mod l2; pub mod l2_u8; pub mod norm_l2; +/// What a per-batch distance kernel yields. +/// +/// The two batch paths produce different shapes: the build-baseline path maps +/// lazily over the batch and allocates nothing, while a `#[target_feature]` +/// kernel must collect eagerly because it cannot be inlined into a lazy +/// closure. A concrete enum keeps both statically dispatched. +/// +/// `Box` would be simpler and is wrong here: hot consumers such +/// as the k-means assignment loop drive these iterators one element at a time +/// through [`crate::kernels::argmin_value_float`], so a trait object turns +/// every `next()` into a virtual call and adds an allocation per batch. +#[cfg(target_arch = "x86_64")] +pub(crate) enum BatchIter { + /// Lazy per-vector map. No allocation. + Lazy(L), + /// Eagerly collected by a `#[target_feature]` kernel. + Eager(std::vec::IntoIter), +} + +#[cfg(target_arch = "x86_64")] +impl> Iterator for BatchIter { + type Item = f32; + + #[inline] + fn next(&mut self) -> Option { + match self { + Self::Lazy(iter) => iter.next(), + Self::Eager(iter) => iter.next(), + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + match self { + Self::Lazy(iter) => iter.size_hint(), + Self::Eager(iter) => iter.size_hint(), + } + } +} + pub use cosine::*; pub use dot::*; pub use hamming::{ diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index d39ee2f0f48..8e410d42580 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -19,6 +19,8 @@ use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use num_traits::{AsPrimitive, Num, real::Real}; use crate::Result; +#[cfg(target_arch = "x86_64")] +use crate::distance::BatchIter; /// Default implementation of dot product. /// @@ -110,6 +112,24 @@ pub fn dot_distance(from: &[T], to: &[T]) -> f32 { pub trait Dot: Num { /// Dot product. fn dot(x: &[Self], y: &[Self]) -> f32; + + /// Dot product of `x` against each `dimension`-sized vector in `batch`. + /// + /// The default calls [`Dot::dot`] per vector. `f32` overrides it so the + /// SIMD tier is chosen once for the whole batch instead of once per + /// vector — on a build whose baseline already implies AVX2, per-vector + /// dispatch costs more than the kernel it selects. + /// + /// Returns `impl Iterator` rather than a trait object: hot consumers drive + /// this one element at a time, so a `Box` would cost a + /// virtual call per element and an allocation per batch. + fn dot_batch<'a>( + x: &'a [Self], + batch: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + batch.chunks_exact(dimension).map(move |y| Self::dot(x, y)) + } } #[cfg(feature = "fp16kernels")] @@ -234,6 +254,83 @@ impl Dot for f32 { // u8 distance kernels in `dot_u8.rs`. dot_f32_dispatched(x, y) } + + fn dot_batch<'a>( + x: &'a [Self], + batch: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + // Exactly one arm compiles. Keeping each a tail expression (rather than + // an early `return` guarded by `cfg`) mirrors `dot_f32_dispatched` and + // avoids an unreachable tail on AVX2-baseline builds. + #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] + { + dot_batch_f32_avx2_baseline(x, batch, dimension) + } + #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] + { + dot_batch_f32_runtime_dispatch(x, batch, dimension) + } + #[cfg(not(target_arch = "x86_64"))] + { + batch.chunks_exact(dimension).map(move |y| Self::dot(x, y)) + } + } +} + +/// AVX2-baseline builds (the default `haswell` wheel): the scalar kernel here +/// inlines and auto-vectorizes exactly as it did before runtime dispatch +/// existed, so a per-vector `*SIMD_SUPPORT` match plus a `#[target_feature]` +/// call boundary would be pure overhead. Dispatch at most once, per batch. +#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] +#[inline] +fn dot_batch_f32_avx2_baseline<'a>( + x: &'a [f32], + batch: &'a [f32], + dimension: usize, +) -> impl Iterator + 'a { + // Only wide vectors benefit from AVX-512: at 16 lanes or fewer a masked + // 512-bit load loses to the plain AVX2 load, and the eager collect adds an + // allocation the baseline path does not need. + if dimension > 16 && matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { + // SAFETY: guarded by the runtime AVX-512 detection above. + return BatchIter::Eager( + unsafe { x86::dot_batch_f32_avx512(x, batch, dimension) }.into_iter(), + ); + } + BatchIter::Lazy( + batch + .chunks_exact(dimension) + .map(move |y| dot_f32_scalar(x, y)), + ) +} + +/// Sub-AVX2 builds: the scalar kernel cannot reach the wide registers, so pick +/// a `#[target_feature]` kernel — once for the batch, not once per vector. +#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[inline] +fn dot_batch_f32_runtime_dispatch<'a>( + x: &'a [f32], + batch: &'a [f32], + dimension: usize, +) -> impl Iterator + 'a { + // SAFETY: each kernel is entered only under its matching runtime detection. + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => { + BatchIter::Eager(unsafe { x86::dot_batch_f32_avx512(x, batch, dimension) }.into_iter()) + } + SimdSupport::Avx2 | SimdSupport::AvxFma => { + BatchIter::Eager(unsafe { x86::dot_batch_f32_avx_fma(x, batch, dimension) }.into_iter()) + } + SimdSupport::Avx => { + BatchIter::Eager(unsafe { x86::dot_batch_f32_avx(x, batch, dimension) }.into_iter()) + } + _ => BatchIter::Lazy( + batch + .chunks_exact(dimension) + .map(move |y| dot_f32_scalar(x, y)), + ), + } } /// Dot product for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 @@ -307,6 +404,52 @@ mod x86 { use crate::simd::x86::hsum256_ps; use crate::simd::{FloatSimd, SIMD}; + /// Dot product of `x` against every `dimension`-sized vector in `batch`, + /// entering the AVX-512 tier once for the whole batch rather than once per + /// vector. + /// + /// # Safety + /// The host must support AVX-512F. + #[target_feature(enable = "avx512f")] + pub(super) unsafe fn dot_batch_f32_avx512( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { dot_f32_avx512(x, y) }) + .collect() + } + + /// As [`dot_batch_f32_avx512`], for the AVX2 and AVX+FMA tiers. + /// + /// # Safety + /// The host must support AVX and FMA. + #[target_feature(enable = "avx,fma")] + pub(super) unsafe fn dot_batch_f32_avx_fma( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { dot_f32_avx_fma(x, y) }) + .collect() + } + + /// As [`dot_batch_f32_avx512`], for the AVX-without-FMA tier. + /// + /// # Safety + /// The host must support AVX. + #[target_feature(enable = "avx")] + pub(super) unsafe fn dot_batch_f32_avx(x: &[f32], batch: &[f32], dimension: usize) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { dot_f32_avx(x, y) }) + .collect() + } + /// AVX-512 path for f64: 8-wide `__m512d` with `vfmadd231pd` per iteration. #[target_feature(enable = "avx512f")] pub unsafe fn dot_f64_avx512(x: &[f64], y: &[f64]) -> f32 { @@ -507,7 +650,7 @@ pub fn dot_distance_batch<'a, T: Dot>( ) -> Box + 'a> { assume_eq!(from.len(), dimension); assume_eq!(to.len() % dimension, 0); - Box::new(to.chunks_exact(dimension).map(|v| dot_distance(from, v))) + Box::new(T::dot_batch(from, to, dimension).map(|d| 1.0 - d)) } fn do_dot_distance_arrow_batch( @@ -531,10 +674,9 @@ where to.value_type() )))?; - let dists = to_values - .as_slice() - .chunks_exact(dimension) - .map(|v| dot_distance(from.as_slice(), v)); + // Route through `dot_distance_batch` rather than mapping `dot_distance` per + // vector, so this entry point gets the same hoisted dispatch. + let dists = dot_distance_batch(from.as_slice(), to_values.as_slice(), dimension); Ok(Arc::new(Float32Array::new( dists.collect(), @@ -849,4 +991,103 @@ mod tests { prop_assert!(approx::relative_eq!(scalar, avx, epsilon = max_error)); } } + + /// `dot_batch` must agree with the per-vector `dot` it replaced, on every + /// build: AVX2-baseline, hoisted-dispatch, and portable fallback all + /// funnel through here. + #[rstest::rstest] + #[case::dim_8(8)] + #[case::dim_16(16)] + #[case::dim_32(32)] + #[case::dim_1024(1024)] + fn test_dot_batch_f32_matches_per_vector_dot(#[case] dimension: usize) { + let num_vectors = 5; + let x: Vec = (0..dimension) + .map(|i| ((i % 13) as f32) * 0.25 + 1.0) + .collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 11) as f32) * 0.5 - 2.0) + .collect(); + + let got: Vec = f32::dot_batch(&x, &batch, dimension).collect(); + let want: Vec = batch + .chunks_exact(dimension) + .map(|y| f32::dot(&x, y)) + .collect(); + + assert_eq!(got.len(), num_vectors); + for (g, w) in got.iter().zip(want.iter()) { + assert!( + approx::relative_eq!(g, w, epsilon = 1e-4), + "dim {dimension}: batch {g} != per-vector {w}" + ); + } + } + + /// `dot_distance_batch` still yields `1.0 - dot`, unchanged by the hoist. + #[test] + fn test_dot_distance_batch_preserves_distance_semantics() { + let dimension = 32; + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.1).collect(); + let batch: Vec = (0..dimension * 3).map(|i| (i as f32) * 0.05).collect(); + + let got: Vec = dot_distance_batch(&x, &batch, dimension).collect(); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + assert!(approx::relative_eq!( + g, + 1.0 - f32::dot(&x, chunk), + epsilon = 1e-5 + )); + } + } + + /// The per-batch `#[target_feature]` kernels are only reached on sub-AVX2 + /// builds or AVX-512 hosts, so call them directly to cover them. + #[cfg(target_arch = "x86_64")] + fn check_dot_batch_kernel(kernel: unsafe fn(&[f32], &[f32], usize) -> Vec) { + for dimension in [8_usize, 16, 40] { + let num_vectors = 3; + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.5 + 1.0).collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 7) as f32) + 1.0) + .collect(); + + let got = unsafe { kernel(&x, &batch, dimension) }; + assert_eq!(got.len(), num_vectors); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + let want = dot_scalar::(&x, chunk); + assert!( + approx::relative_eq!(g, want, epsilon = 1e-4), + "dim {dimension}: kernel {g} != scalar {want}" + ); + } + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_batch_avx_fma_matches_scalar() { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx_fma); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_batch_avx_matches_scalar() { + if !std::is_x86_feature_detected!("avx") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn test_dot_batch_avx512_matches_scalar() { + if !std::is_x86_feature_detected!("avx512f") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx512); + } } diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 40687e33f85..05c88f973c7 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -26,14 +26,32 @@ use lance_core::utils::cpu::SIMD_SUPPORT; use lance_core::utils::cpu::SimdSupport; use num_traits::{AsPrimitive, Num}; +#[cfg(target_arch = "x86_64")] +use crate::distance::BatchIter; + /// Calculate the L2 distance between two vectors. /// pub trait L2: Num { /// Calculate the L2 distance between two vectors. fn l2(x: &[Self], y: &[Self]) -> f32; - fn l2_batch(x: &[Self], y: &[Self], dimension: usize) -> impl Iterator { - y.chunks_exact(dimension).map(|v| Self::l2(x, v)) + /// L2 distance from `x` to each `dimension`-sized vector in `y`. + /// + /// The default calls [`L2::l2`] per vector. `f32` overrides it so the SIMD + /// tier is chosen once for the whole batch instead of once per vector — + /// on a build whose baseline already implies AVX2, per-vector dispatch + /// costs more than the kernel it selects. + /// + /// Returns `impl Iterator` rather than a trait object: the k-means + /// assignment loop drives this one element at a time, so a + /// `Box` would cost a virtual call per element and an + /// allocation per batch. + fn l2_batch<'a>( + x: &'a [Self], + y: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + y.chunks_exact(dimension).map(move |v| Self::l2(x, v)) } } @@ -265,6 +283,69 @@ impl L2 for f32 { // scalar fallback. l2_f32_dispatched(x, y) } + + fn l2_batch<'a>( + x: &'a [Self], + y: &'a [Self], + dimension: usize, + ) -> impl Iterator + 'a { + // Exactly one arm compiles; see `Dot::dot_batch` for f32. + #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] + { + l2_batch_f32_avx2_baseline(x, y, dimension) + } + #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] + { + l2_batch_f32_runtime_dispatch(x, y, dimension) + } + #[cfg(not(target_arch = "x86_64"))] + { + y.chunks_exact(dimension).map(move |v| Self::l2(x, v)) + } + } +} + +/// AVX2-baseline builds: the scalar kernel here inlines and auto-vectorizes as +/// it did before runtime dispatch existed, so per-vector dispatch would be pure +/// overhead. Dispatch at most once, per batch. See `Dot::dot_batch` for f32. +#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] +#[inline] +fn l2_batch_f32_avx2_baseline<'a>( + x: &'a [f32], + y: &'a [f32], + dimension: usize, +) -> impl Iterator + 'a { + // Only wide vectors benefit from AVX-512: at 16 lanes or fewer a masked + // 512-bit load loses to the plain AVX2 load, and the eager collect adds an + // allocation the baseline path does not need. + if dimension > 16 && matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { + // SAFETY: guarded by the runtime AVX-512 detection above. + return BatchIter::Eager(unsafe { x86::l2_batch_f32_avx512(x, y, dimension) }.into_iter()); + } + BatchIter::Lazy(y.chunks_exact(dimension).map(move |v| l2_f32_scalar(x, v))) +} + +/// Sub-AVX2 builds: pick a `#[target_feature]` kernel once for the batch. +#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[inline] +fn l2_batch_f32_runtime_dispatch<'a>( + x: &'a [f32], + y: &'a [f32], + dimension: usize, +) -> impl Iterator + 'a { + // SAFETY: each kernel is entered only under its matching runtime detection. + match *SIMD_SUPPORT { + SimdSupport::Avx512 | SimdSupport::Avx512FP16 => { + BatchIter::Eager(unsafe { x86::l2_batch_f32_avx512(x, y, dimension) }.into_iter()) + } + SimdSupport::Avx2 | SimdSupport::AvxFma => { + BatchIter::Eager(unsafe { x86::l2_batch_f32_avx_fma(x, y, dimension) }.into_iter()) + } + SimdSupport::Avx => { + BatchIter::Eager(unsafe { x86::l2_batch_f32_avx(x, y, dimension) }.into_iter()) + } + _ => BatchIter::Lazy(y.chunks_exact(dimension).map(move |v| l2_f32_scalar(x, v))), + } } /// L2 distance for f32, runtime-dispatched via `SIMD_SUPPORT` on x86_64 @@ -346,6 +427,52 @@ mod x86 { use crate::simd::x86::hsum256_ps; use crate::simd::{FloatSimd, SIMD}; + /// L2 distance from `x` to every `dimension`-sized vector in `batch`, with + /// the AVX-512 tier entered once for the whole batch rather than once per + /// vector. + /// + /// # Safety + /// The host must support AVX-512F. + #[target_feature(enable = "avx512f")] + pub(super) unsafe fn l2_batch_f32_avx512( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { l2_f32_avx512(x, y) }) + .collect() + } + + /// As [`l2_batch_f32_avx512`], for the AVX+FMA and AVX2 tiers. + /// + /// # Safety + /// The host must support AVX and FMA. + #[target_feature(enable = "avx,fma")] + pub(super) unsafe fn l2_batch_f32_avx_fma( + x: &[f32], + batch: &[f32], + dimension: usize, + ) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { l2_f32_avx_fma(x, y) }) + .collect() + } + + /// As [`l2_batch_f32_avx512`], for the AVX-without-FMA tier. + /// + /// # Safety + /// The host must support AVX. + #[target_feature(enable = "avx")] + pub(super) unsafe fn l2_batch_f32_avx(x: &[f32], batch: &[f32], dimension: usize) -> Vec { + batch + .chunks_exact(dimension) + .map(|y| unsafe { l2_f32_avx(x, y) }) + .collect() + } + /// AVX-512 path for f64: 8-wide `__m512d` with `vsubpd` + `vfmadd231pd` per iteration. #[target_feature(enable = "avx512f")] pub unsafe fn l2_f64_avx512(x: &[f64], y: &[f64]) -> f32 { @@ -1171,4 +1298,86 @@ mod tests { assert_relative_eq!(d1[0], 0.0); // q1 == target[0] assert_relative_eq!(d2[1], 0.0); // q2 == target[1] } + + /// `l2_batch` must agree with the per-vector `l2` it replaced, on every + /// build: the AVX2-baseline path, the hoisted-dispatch path, and the + /// portable fallback all funnel through here. + #[rstest::rstest] + #[case::dim_8(8)] + #[case::dim_16(16)] + #[case::dim_32(32)] + #[case::dim_1024(1024)] + fn test_l2_batch_f32_matches_per_vector_l2(#[case] dimension: usize) { + let num_vectors = 5; + let x: Vec = (0..dimension) + .map(|i| ((i % 13) as f32) * 0.25 + 1.0) + .collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 11) as f32) * 0.5 - 2.0) + .collect(); + + let got: Vec = f32::l2_batch(&x, &batch, dimension).collect(); + let want: Vec = batch + .chunks_exact(dimension) + .map(|y| f32::l2(&x, y)) + .collect(); + + assert_eq!(got.len(), num_vectors); + for (g, w) in got.iter().zip(want.iter()) { + assert!( + approx::relative_eq!(g, w, epsilon = 1e-4), + "dim {dimension}: batch {g} != per-vector {w}" + ); + } + } + + /// The per-batch `#[target_feature]` kernels are only reached on sub-AVX2 + /// builds or AVX-512 hosts, so call them directly to cover them. + #[cfg(target_arch = "x86_64")] + fn check_l2_batch_kernel(kernel: unsafe fn(&[f32], &[f32], usize) -> Vec) { + for dimension in [8_usize, 16, 40] { + let num_vectors = 3; + let x: Vec = (0..dimension).map(|i| (i as f32) * 0.5 + 1.0).collect(); + let batch: Vec = (0..dimension * num_vectors) + .map(|i| ((i % 7) as f32) + 1.0) + .collect(); + + let got = unsafe { kernel(&x, &batch, dimension) }; + assert_eq!(got.len(), num_vectors); + for (chunk, &g) in batch.chunks_exact(dimension).zip(got.iter()) { + let want = l2_scalar::(&x, chunk); + assert!( + approx::relative_eq!(g, want, epsilon = 1e-4), + "dim {dimension}: kernel {g} != scalar {want}" + ); + } + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_batch_avx_fma_matches_scalar() { + if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx_fma); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_batch_avx_matches_scalar() { + if !std::is_x86_feature_detected!("avx") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn test_l2_batch_avx512_matches_scalar() { + if !std::is_x86_feature_detected!("avx512f") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx512); + } } From 97865b52362849b2ef7f96a502dbf96d8c69ffba Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 20:48:12 -0400 Subject: [PATCH 11/15] perf: keep the batch iterator wrapper-free on AVX2-baseline builds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Benchmarking the previous commit on two hosts turned up a regression it had introduced: L2(simd,f32x8) was +154% on an AVX2-only EPYC 7B13 and +78.7% on an AVX-512 Xeon, while the dim-1024 benchmarks were flat. The cause was the `BatchIter` enum. Wrapping the batch iterator loses two things the bare `Map` has, and both matter only when the per-vector work is small: - `Map` overrides `Iterator::fold` to drive `ChunksExact` in one inlined, auto-vectorized loop. Through the enum, `fold` fell back to calling `next()` per element. - `Map` is `TrustedLen`, so `.collect()` preallocates exactly. No enum can be, and `TrustedLen` cannot be implemented outside std. `fold` can be delegated; `TrustedLen` cannot, and the hot consumers do not all use `fold` — `argmin_value_float` drives the iterator through `enumerate()`, and the arrow batch paths `.collect()`. So on an AVX2-baseline build, return the bare `Map` over the inlined scalar kernel: byte-for-byte the iterator the code produced before runtime dispatch existed. That removes the per-vector dispatch without adding anything, which is the only shape that is non-regressing by construction rather than by measurement. `BatchIter` now exists only for sub-AVX2 builds, where the eager `#[target_feature]` kernels need a common return type and there is no prior behaviour to regress. It delegates `fold`/`for_each` and implements `ExactSizeIterator` so that path stays reasonable too. --- rust/lance-linalg/src/distance.rs | 53 ++++++++++++++++++++++++--- rust/lance-linalg/src/distance/dot.rs | 39 +++++--------------- rust/lance-linalg/src/distance/l2.rs | 26 ++----------- 3 files changed, 61 insertions(+), 57 deletions(-) diff --git a/rust/lance-linalg/src/distance.rs b/rust/lance-linalg/src/distance.rs index 0a2cb31a7d4..42977c56d2c 100644 --- a/rust/lance-linalg/src/distance.rs +++ b/rust/lance-linalg/src/distance.rs @@ -32,11 +32,12 @@ pub mod norm_l2; /// kernel must collect eagerly because it cannot be inlined into a lazy /// closure. A concrete enum keeps both statically dispatched. /// -/// `Box` would be simpler and is wrong here: hot consumers such -/// as the k-means assignment loop drive these iterators one element at a time -/// through [`crate::kernels::argmin_value_float`], so a trait object turns -/// every `next()` into a virtual call and adds an allocation per batch. -#[cfg(target_arch = "x86_64")] +/// Only sub-AVX2 builds need this. On an AVX2-baseline build the batch methods +/// return the bare `Map` instead, because any wrapper — trait object or enum — +/// loses `TrustedLen` (so `.collect()` stops preallocating) and loses +/// `Map::fold`'s inlined loop. Benchmarks showed that costing 2.5x on the dim-8 +/// batch, far more than the per-vector dispatch it was meant to remove. +#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] pub(crate) enum BatchIter { /// Lazy per-vector map. No allocation. Lazy(L), @@ -44,7 +45,7 @@ pub(crate) enum BatchIter { Eager(std::vec::IntoIter), } -#[cfg(target_arch = "x86_64")] +#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] impl> Iterator for BatchIter { type Item = f32; @@ -63,6 +64,46 @@ impl> Iterator for BatchIter { Self::Eager(iter) => iter.size_hint(), } } + + /// Delegated, not defaulted. `Map` overrides `fold` to drive the underlying + /// `ChunksExact` in one inlined, auto-vectorized loop; the default `fold` + /// would instead call `next()` per element, paying an enum branch and + /// losing that loop. On the dim-8 batch that costs ~2.5x. + #[inline] + fn fold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + match self { + Self::Lazy(iter) => iter.fold(init, f), + Self::Eager(iter) => iter.fold(init, f), + } + } + + /// `for_each`, `sum` and `collect` all route through `fold`, so delegating + /// it covers them too. (`try_fold` cannot be overridden on stable: its + /// `Try` bound is unstable.) + #[inline] + fn for_each(self, f: F) + where + F: FnMut(Self::Item), + { + match self { + Self::Lazy(iter) => iter.for_each(f), + Self::Eager(iter) => iter.for_each(f), + } + } +} + +#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +impl> ExactSizeIterator for BatchIter { + #[inline] + fn len(&self) -> usize { + match self { + Self::Lazy(iter) => iter.len(), + Self::Eager(iter) => iter.len(), + } + } } pub use cosine::*; diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 8e410d42580..e4f558ea290 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -19,7 +19,7 @@ use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use num_traits::{AsPrimitive, Num, real::Real}; use crate::Result; -#[cfg(target_arch = "x86_64")] +#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] use crate::distance::BatchIter; /// Default implementation of dot product. @@ -263,9 +263,17 @@ impl Dot for f32 { // Exactly one arm compiles. Keeping each a tail expression (rather than // an early `return` guarded by `cfg`) mirrors `dot_f32_dispatched` and // avoids an unreachable tail on AVX2-baseline builds. + // AVX2-baseline build (the default `haswell` wheel): return exactly what + // the code returned before runtime dispatch existed — a lazy `Map` over + // the inlined, auto-vectorized scalar kernel. Wrapping it in anything + // (a trait object, or an enum) costs more than the dispatch it removes: + // `Map` is `TrustedLen`, so `.collect()` preallocates, + // and `Map::fold` drives `ChunksExact` in one inlined loop. #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] { - dot_batch_f32_avx2_baseline(x, batch, dimension) + batch + .chunks_exact(dimension) + .map(move |y| dot_f32_scalar(x, y)) } #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] { @@ -278,33 +286,6 @@ impl Dot for f32 { } } -/// AVX2-baseline builds (the default `haswell` wheel): the scalar kernel here -/// inlines and auto-vectorizes exactly as it did before runtime dispatch -/// existed, so a per-vector `*SIMD_SUPPORT` match plus a `#[target_feature]` -/// call boundary would be pure overhead. Dispatch at most once, per batch. -#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] -#[inline] -fn dot_batch_f32_avx2_baseline<'a>( - x: &'a [f32], - batch: &'a [f32], - dimension: usize, -) -> impl Iterator + 'a { - // Only wide vectors benefit from AVX-512: at 16 lanes or fewer a masked - // 512-bit load loses to the plain AVX2 load, and the eager collect adds an - // allocation the baseline path does not need. - if dimension > 16 && matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { - // SAFETY: guarded by the runtime AVX-512 detection above. - return BatchIter::Eager( - unsafe { x86::dot_batch_f32_avx512(x, batch, dimension) }.into_iter(), - ); - } - BatchIter::Lazy( - batch - .chunks_exact(dimension) - .map(move |y| dot_f32_scalar(x, y)), - ) -} - /// Sub-AVX2 builds: the scalar kernel cannot reach the wide registers, so pick /// a `#[target_feature]` kernel — once for the batch, not once per vector. #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 05c88f973c7..462d3aed26e 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -26,7 +26,7 @@ use lance_core::utils::cpu::SIMD_SUPPORT; use lance_core::utils::cpu::SimdSupport; use num_traits::{AsPrimitive, Num}; -#[cfg(target_arch = "x86_64")] +#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] use crate::distance::BatchIter; /// Calculate the L2 distance between two vectors. @@ -290,9 +290,11 @@ impl L2 for f32 { dimension: usize, ) -> impl Iterator + 'a { // Exactly one arm compiles; see `Dot::dot_batch` for f32. + // See `Dot::dot_batch` for f32: on an AVX2-baseline build this is exactly + // the pre-dispatch iterator, wrapper-free. #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] { - l2_batch_f32_avx2_baseline(x, y, dimension) + y.chunks_exact(dimension).map(move |v| l2_f32_scalar(x, v)) } #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] { @@ -305,26 +307,6 @@ impl L2 for f32 { } } -/// AVX2-baseline builds: the scalar kernel here inlines and auto-vectorizes as -/// it did before runtime dispatch existed, so per-vector dispatch would be pure -/// overhead. Dispatch at most once, per batch. See `Dot::dot_batch` for f32. -#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] -#[inline] -fn l2_batch_f32_avx2_baseline<'a>( - x: &'a [f32], - y: &'a [f32], - dimension: usize, -) -> impl Iterator + 'a { - // Only wide vectors benefit from AVX-512: at 16 lanes or fewer a masked - // 512-bit load loses to the plain AVX2 load, and the eager collect adds an - // allocation the baseline path does not need. - if dimension > 16 && matches!(*SIMD_SUPPORT, SimdSupport::Avx512 | SimdSupport::Avx512FP16) { - // SAFETY: guarded by the runtime AVX-512 detection above. - return BatchIter::Eager(unsafe { x86::l2_batch_f32_avx512(x, y, dimension) }.into_iter()); - } - BatchIter::Lazy(y.chunks_exact(dimension).map(move |v| l2_f32_scalar(x, v))) -} - /// Sub-AVX2 builds: pick a `#[target_feature]` kernel once for the batch. #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] #[inline] From 334229aab7d1b37d0b7c7404cd6283234c69f990 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 21:43:36 -0400 Subject: [PATCH 12/15] perf: keep the AVX+FMA kernel on the AVX2-baseline batch path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Hoisting the tier choice out of the per-vector loop is only half the job. The first attempt also replaced the kernel with the auto-vectorized scalar one, on the theory that an AVX2-baseline build makes the SIMD kernel redundant. Benchmarks say otherwise: at dim 8, `L2(simd,f32x8)` went from -41.6% (before this work) to +156% against upstream/main. Removing a cheap branch is worthless if it costs the kernel behind it. The baseline already guarantees avx2+fma, so call the AVX+FMA kernel directly: no runtime check, and the `#[target_feature]` contract is satisfied statically, so it inlines into the loop. Measured on an AMD EPYC 7B13 (avx2, fma, no avx512f), against upstream/main: before this work with this commit L2(simd,f32x8) dim 8 -41.6% -45.4% L2(f32, simd) dim 1024 +1.5% +0.4% Dot(f32, SIMD) -0.1% -0.5% The iterator stays a bare `Map`: `Map` is `TrustedLen`, so `.collect()` preallocates, and `Map::fold` drives `ChunksExact` in one inlined loop. Any wrapper — trait object or enum — loses both. --- rust/lance-linalg/src/distance/dot.rs | 21 ++++++++++++++------- rust/lance-linalg/src/distance/l2.rs | 10 +++++++--- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index e4f558ea290..00ea37d01ed 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -263,17 +263,24 @@ impl Dot for f32 { // Exactly one arm compiles. Keeping each a tail expression (rather than // an early `return` guarded by `cfg`) mirrors `dot_f32_dispatched` and // avoids an unreachable tail on AVX2-baseline builds. - // AVX2-baseline build (the default `haswell` wheel): return exactly what - // the code returned before runtime dispatch existed — a lazy `Map` over - // the inlined, auto-vectorized scalar kernel. Wrapping it in anything - // (a trait object, or an enum) costs more than the dispatch it removes: - // `Map` is `TrustedLen`, so `.collect()` preallocates, - // and `Map::fold` drives `ChunksExact` in one inlined loop. + // AVX2-baseline build (the default `haswell` wheel). Hoist the tier + // choice out of the loop, but keep the SIMD kernel: the baseline already + // guarantees avx2+fma, so call the AVX+FMA kernel directly rather than + // re-checking per vector. Falling back to the scalar kernel here would + // lose ~4x at small dimensions, which is where batch calls live (PQ + // sub-vectors are 8 wide). + // + // The iterator is a bare `Map`: `Map` is `TrustedLen`, + // so `.collect()` preallocates, and `Map::fold` drives `ChunksExact` in + // one inlined loop. Any wrapper — trait object or enum — loses both. #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] { + // SAFETY: avx2+fma are enabled for the whole crate by the build + // baseline, so the kernel's `#[target_feature]` contract holds + // statically and it inlines here. batch .chunks_exact(dimension) - .map(move |y| dot_f32_scalar(x, y)) + .map(move |y| unsafe { x86::dot_f32_avx_fma(x, y) }) } #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] { diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 462d3aed26e..bc71eb8a502 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -290,11 +290,15 @@ impl L2 for f32 { dimension: usize, ) -> impl Iterator + 'a { // Exactly one arm compiles; see `Dot::dot_batch` for f32. - // See `Dot::dot_batch` for f32: on an AVX2-baseline build this is exactly - // the pre-dispatch iterator, wrapper-free. + // See `Dot::dot_batch` for f32. #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] { - y.chunks_exact(dimension).map(move |v| l2_f32_scalar(x, v)) + // SAFETY: the build baseline enables avx2+fma, which imply avx+fma, + // so the kernel's `#[target_feature]` contract is met statically and + // it inlines into this loop. No runtime check, and no falling back + // to the scalar kernel — at small dimensions that costs ~4x. + y.chunks_exact(dimension) + .map(move |v| unsafe { x86::l2_f32_avx_fma(x, v) }) } #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] { From 07cd9670ea1465c540d2c024fcc1e1642f262459 Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Thu, 9 Jul 2026 22:47:36 -0400 Subject: [PATCH 13/15] fix: require avx2 and fma before taking the direct-kernel batch path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `cfg(target_feature = "avx2")` does not imply `fma` — `-C target-feature=+avx2` sets one and not the other. The batch path guarded on it calls a kernel declared `#[target_feature(enable = "avx,fma")]`, so such a build emitted `vfmadd` while claiming a baseline without FMA. Gate the direct-kernel arm on both features, and make the runtime-dispatch arm its exact complement so the two remain total. Verified by building the crate at `+avx2`, `+avx2,+fma`, `+avx`, `target-cpu=haswell`, and the default baseline. --- rust/lance-linalg/src/distance.rs | 15 ++++++++++++--- rust/lance-linalg/src/distance/dot.rs | 21 +++++++++++++++++---- rust/lance-linalg/src/distance/l2.rs | 21 +++++++++++++++++---- 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/rust/lance-linalg/src/distance.rs b/rust/lance-linalg/src/distance.rs index 42977c56d2c..d5f0dcba7aa 100644 --- a/rust/lance-linalg/src/distance.rs +++ b/rust/lance-linalg/src/distance.rs @@ -37,7 +37,10 @@ pub mod norm_l2; /// loses `TrustedLen` (so `.collect()` stops preallocating) and loses /// `Map::fold`'s inlined loop. Benchmarks showed that costing 2.5x on the dim-8 /// batch, far more than the per-vector dispatch it was meant to remove. -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] pub(crate) enum BatchIter { /// Lazy per-vector map. No allocation. Lazy(L), @@ -45,7 +48,10 @@ pub(crate) enum BatchIter { Eager(std::vec::IntoIter), } -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] impl> Iterator for BatchIter { type Item = f32; @@ -95,7 +101,10 @@ impl> Iterator for BatchIter { } } -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] impl> ExactSizeIterator for BatchIter { #[inline] fn len(&self) -> usize { diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 00ea37d01ed..ac138207e46 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -19,7 +19,10 @@ use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport}; use num_traits::{AsPrimitive, Num, real::Real}; use crate::Result; -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] use crate::distance::BatchIter; /// Default implementation of dot product. @@ -273,7 +276,11 @@ impl Dot for f32 { // The iterator is a bare `Map`: `Map` is `TrustedLen`, // so `.collect()` preallocates, and `Map::fold` drives `ChunksExact` in // one inlined loop. Any wrapper — trait object or enum — loses both. - #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx2", + target_feature = "fma" + ))] { // SAFETY: avx2+fma are enabled for the whole crate by the build // baseline, so the kernel's `#[target_feature]` contract holds @@ -282,7 +289,10 @@ impl Dot for f32 { .chunks_exact(dimension) .map(move |y| unsafe { x86::dot_f32_avx_fma(x, y) }) } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] { dot_batch_f32_runtime_dispatch(x, batch, dimension) } @@ -295,7 +305,10 @@ impl Dot for f32 { /// Sub-AVX2 builds: the scalar kernel cannot reach the wide registers, so pick /// a `#[target_feature]` kernel — once for the batch, not once per vector. -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] #[inline] fn dot_batch_f32_runtime_dispatch<'a>( x: &'a [f32], diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index bc71eb8a502..56b9ff0c432 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -26,7 +26,10 @@ use lance_core::utils::cpu::SIMD_SUPPORT; use lance_core::utils::cpu::SimdSupport; use num_traits::{AsPrimitive, Num}; -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] use crate::distance::BatchIter; /// Calculate the L2 distance between two vectors. @@ -291,7 +294,11 @@ impl L2 for f32 { ) -> impl Iterator + 'a { // Exactly one arm compiles; see `Dot::dot_batch` for f32. // See `Dot::dot_batch` for f32. - #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx2", + target_feature = "fma" + ))] { // SAFETY: the build baseline enables avx2+fma, which imply avx+fma, // so the kernel's `#[target_feature]` contract is met statically and @@ -300,7 +307,10 @@ impl L2 for f32 { y.chunks_exact(dimension) .map(move |v| unsafe { x86::l2_f32_avx_fma(x, v) }) } - #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] { l2_batch_f32_runtime_dispatch(x, y, dimension) } @@ -312,7 +322,10 @@ impl L2 for f32 { } /// Sub-AVX2 builds: pick a `#[target_feature]` kernel once for the batch. -#[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))] +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] #[inline] fn l2_batch_f32_runtime_dispatch<'a>( x: &'a [f32], From 6b94ccdbe3c4e60d4f62024bd87c41797505805c Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Fri, 10 Jul 2026 00:34:05 -0400 Subject: [PATCH 14/15] perf: only use the AVX kernel where the scalar path degenerates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `dot_scalar`/`l2_scalar` chunk by 16 lanes. At or below that width the chunking degenerates to the scalar remainder loop and vectorizes nothing, which is why the explicit AVX kernel is worth ~40% at dim 8 — the width PQ sub-vectors use. Above 16 lanes the autovectorizer is already good, and on Zen 3 the 8-wide kernel loses to it. Measured on an AMD EPYC 7B13 (avx2, fma, no avx512f), two repetitions each, against upstream/main: kernel everywhere this commit L2(simd,f32x8) -42.2% / -41.4% (unchanged, kernel) Dot(f32, SIMD) +1.9% / +3.6% (scalar, == base) So above the threshold keep exactly the kernel the pre-dispatch code used. The wide path is then byte-identical to base and cannot regress; the narrow path keeps the win. The branch is loop-invariant. --- rust/lance-linalg/src/distance/dot.rs | 18 ++++++++++++++---- rust/lance-linalg/src/distance/l2.rs | 21 ++++++++++++++++----- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index ac138207e46..76837ffcaa9 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -282,12 +282,22 @@ impl Dot for f32 { target_feature = "fma" ))] { + // See `L2::l2_batch` for f32: below 16 lanes `dot_scalar`'s chunking + // degenerates to a scalar remainder loop, so the explicit AVX kernel + // wins big; above it the autovectorizer is already good and the + // 8-wide kernel can lose, so keep the pre-dispatch kernel exactly. + // // SAFETY: avx2+fma are enabled for the whole crate by the build // baseline, so the kernel's `#[target_feature]` contract holds - // statically and it inlines here. - batch - .chunks_exact(dimension) - .map(move |y| unsafe { x86::dot_f32_avx_fma(x, y) }) + // statically. + let narrow = dimension <= 16; + batch.chunks_exact(dimension).map(move |y| { + if narrow { + unsafe { x86::dot_f32_avx_fma(x, y) } + } else { + dot_f32_scalar(x, y) + } + }) } #[cfg(all( target_arch = "x86_64", diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 56b9ff0c432..41790d8df4a 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -300,12 +300,23 @@ impl L2 for f32 { target_feature = "fma" ))] { + // `l2_scalar::<_, _, 16>` chunks the vector by 16 lanes. At or below + // that width the chunking degenerates to its scalar remainder loop + // and vectorizes nothing, so the explicit AVX kernel is worth ~40%. + // Above it the autovectorizer already does well and the 8-wide + // kernel can lose, so keep the exact kernel the pre-dispatch code + // used and stay non-regressing by construction. + // // SAFETY: the build baseline enables avx2+fma, which imply avx+fma, - // so the kernel's `#[target_feature]` contract is met statically and - // it inlines into this loop. No runtime check, and no falling back - // to the scalar kernel — at small dimensions that costs ~4x. - y.chunks_exact(dimension) - .map(move |v| unsafe { x86::l2_f32_avx_fma(x, v) }) + // so the kernel's `#[target_feature]` contract is met statically. + let narrow = dimension <= 16; + y.chunks_exact(dimension).map(move |v| { + if narrow { + unsafe { x86::l2_f32_avx_fma(x, v) } + } else { + l2_f32_scalar(x, v) + } + }) } #[cfg(all( target_arch = "x86_64", From b1670de268cf6cf67fdac8cfe3a4a6cd194dcaef Mon Sep 17 00:00:00 2001 From: tobocop2 <5562156+tobocop2@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:16:53 -0400 Subject: [PATCH 15/15] fix: gate f32 batch dispatch kernels to sub-avx2+fma builds The per-batch `dot_batch_f32_avx*` / `l2_batch_f32_avx*` wrappers are only called from the runtime-dispatch path, which is compiled out when the build baseline already guarantees avx2+fma (the repo's `target-cpu=haswell` default). Under that baseline the wrappers, their test helper, and their tests were unused, so `cargo clippy -D warnings` failed on dead_code. Gate the definitions and their tests with the same `not(all(target_feature = "avx2", target_feature = "fma"))` cfg as their caller, so definition, dispatch, and test compile together. No behavior change: at the avx2+fma baseline `dot_batch`/`l2_batch` already inline the kernel directly; below it the runtime dispatch and these wrappers remain. --- rust/lance-linalg/src/distance/dot.rs | 29 +++++++++++++++++++++++---- rust/lance-linalg/src/distance/l2.rs | 29 +++++++++++++++++++++++---- 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 76837ffcaa9..3411d7e8dce 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -421,6 +421,11 @@ mod x86 { /// /// # Safety /// The host must support AVX-512F. + /// + /// Only compiled for builds whose baseline is below avx2+fma; at or above + /// that baseline `dot_batch` inlines the kernel directly and never runtime- + /// dispatches, so this wrapper would be dead code (see `dot_batch`). + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] #[target_feature(enable = "avx512f")] pub(super) unsafe fn dot_batch_f32_avx512( x: &[f32], @@ -437,6 +442,7 @@ mod x86 { /// /// # Safety /// The host must support AVX and FMA. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] #[target_feature(enable = "avx,fma")] pub(super) unsafe fn dot_batch_f32_avx_fma( x: &[f32], @@ -453,6 +459,7 @@ mod x86 { /// /// # Safety /// The host must support AVX. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] #[target_feature(enable = "avx")] pub(super) unsafe fn dot_batch_f32_avx(x: &[f32], batch: &[f32], dimension: usize) -> Vec { batch @@ -1054,7 +1061,10 @@ mod tests { /// The per-batch `#[target_feature]` kernels are only reached on sub-AVX2 /// builds or AVX-512 hosts, so call them directly to cover them. - #[cfg(target_arch = "x86_64")] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] fn check_dot_batch_kernel(kernel: unsafe fn(&[f32], &[f32], usize) -> Vec) { for dimension in [8_usize, 16, 40] { let num_vectors = 3; @@ -1075,7 +1085,12 @@ mod tests { } } - #[cfg(target_arch = "x86_64")] + // The runtime-dispatch batch kernels only exist in sub-avx2+fma builds + // (see `x86::dot_batch_f32_avx512`), so gate their tests the same way. + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] #[test] fn test_dot_batch_avx_fma_matches_scalar() { if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { @@ -1084,7 +1099,10 @@ mod tests { check_dot_batch_kernel(x86::dot_batch_f32_avx_fma); } - #[cfg(target_arch = "x86_64")] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] #[test] fn test_dot_batch_avx_matches_scalar() { if !std::is_x86_feature_detected!("avx") { @@ -1093,7 +1111,10 @@ mod tests { check_dot_batch_kernel(x86::dot_batch_f32_avx); } - #[cfg(target_arch = "x86_64")] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] #[test] fn test_dot_batch_avx512_matches_scalar() { if !std::is_x86_feature_detected!("avx512f") { diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 41790d8df4a..7ee2f4bc537 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -443,6 +443,11 @@ mod x86 { /// /// # Safety /// The host must support AVX-512F. + /// + /// Only compiled for builds whose baseline is below avx2+fma; at or above + /// that baseline `l2_batch` inlines the kernel directly and never runtime- + /// dispatches, so this wrapper would be dead code (see `l2_batch`). + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] #[target_feature(enable = "avx512f")] pub(super) unsafe fn l2_batch_f32_avx512( x: &[f32], @@ -459,6 +464,7 @@ mod x86 { /// /// # Safety /// The host must support AVX and FMA. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] #[target_feature(enable = "avx,fma")] pub(super) unsafe fn l2_batch_f32_avx_fma( x: &[f32], @@ -475,6 +481,7 @@ mod x86 { /// /// # Safety /// The host must support AVX. + #[cfg(not(all(target_feature = "avx2", target_feature = "fma")))] #[target_feature(enable = "avx")] pub(super) unsafe fn l2_batch_f32_avx(x: &[f32], batch: &[f32], dimension: usize) -> Vec { batch @@ -1343,7 +1350,10 @@ mod tests { /// The per-batch `#[target_feature]` kernels are only reached on sub-AVX2 /// builds or AVX-512 hosts, so call them directly to cover them. - #[cfg(target_arch = "x86_64")] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] fn check_l2_batch_kernel(kernel: unsafe fn(&[f32], &[f32], usize) -> Vec) { for dimension in [8_usize, 16, 40] { let num_vectors = 3; @@ -1364,7 +1374,12 @@ mod tests { } } - #[cfg(target_arch = "x86_64")] + // The runtime-dispatch batch kernels only exist in sub-avx2+fma builds + // (see `x86::l2_batch_f32_avx512`), so gate their tests the same way. + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] #[test] fn test_l2_batch_avx_fma_matches_scalar() { if !std::is_x86_feature_detected!("avx") || !std::is_x86_feature_detected!("fma") { @@ -1373,7 +1388,10 @@ mod tests { check_l2_batch_kernel(x86::l2_batch_f32_avx_fma); } - #[cfg(target_arch = "x86_64")] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] #[test] fn test_l2_batch_avx_matches_scalar() { if !std::is_x86_feature_detected!("avx") { @@ -1382,7 +1400,10 @@ mod tests { check_l2_batch_kernel(x86::l2_batch_f32_avx); } - #[cfg(target_arch = "x86_64")] + #[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] #[test] fn test_l2_batch_avx512_matches_scalar() { if !std::is_x86_feature_detected!("avx512f") {