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..3940ed7b683 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 `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/Cargo.lock b/Cargo.lock index d8dcf1fd67a..586cce097f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4932,6 +4932,7 @@ dependencies = [ "proptest", "rand 0.9.4", "rayon", + "rstest", ] [[package]] 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 860590847e6..c5631d26f10 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -329,6 +329,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))?; @@ -347,6 +348,37 @@ 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) +/// } +/// +/// 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(); + 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..c4d5a976cbb 100644 --- a/rust/lance-core/src/utils/cpu.rs +++ b/rust/lance-core/src/utils/cpu.rs @@ -1,14 +1,29 @@ // 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 + 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, @@ -16,6 +31,132 @@ 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. Obtain one with [`simd_info()`]. +#[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. 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, + 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::with_capacity(17); + 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")))] @@ -42,8 +183,19 @@ 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. + 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 +290,68 @@ mod aarch64 { false } } + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[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()); + } + + /// 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")] + #[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/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/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.rs b/rust/lance-linalg/src/distance.rs index 23d1cae2d63..d5f0dcba7aa 100644 --- a/rust/lance-linalg/src/distance.rs +++ b/rust/lance-linalg/src/distance.rs @@ -25,6 +25,96 @@ 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. +/// +/// 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(all(target_feature = "avx2", target_feature = "fma")) +))] +pub(crate) enum BatchIter { + /// Lazy per-vector map. No allocation. + Lazy(L), + /// Eagerly collected by a `#[target_feature]` kernel. + Eager(std::vec::IntoIter), +} + +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +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(), + } + } + + /// 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(all(target_feature = "avx2", target_feature = "fma")) +))] +impl> ExactSizeIterator for BatchIter { + #[inline] + fn len(&self) -> usize { + match self { + Self::Lazy(iter) => iter.len(), + Self::Eager(iter) => iter.len(), + } + } +} + pub use cosine::*; pub use dot::*; pub use hamming::{ diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index 995191b77eb..457c0c0dd6a 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,88 +183,348 @@ 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>( + 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) + } + } + + #[inline] + 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), + } + } + #[cfg(not(target_arch = "x86_64"))] + { + cosine_once_16_other(x, x_norm, y) + } + } + + /// Portable scalar `cosine_once` for length-8 vectors. Matches the SIMD + /// path modulo summation order. + #[cfg(target_arch = "x86_64")] + #[inline] + 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]; + } + 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]; + } + 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() + } + + /// 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, - 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() + 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(), + } } } -impl Cosine for f32 { +/// 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] - 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 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() } - 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); - } + } + + #[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() } - 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() } #[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 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 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); + 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() } - let xy = xy16.reduce_sum() + xy8.reduce_sum() + dot(&x[aligned_len..], &y[aligned_len..]); - 1.0 - xy / x_norm / y_norm + } +} + +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) + } + + #[allow(unreachable_code)] fn cosine_batch<'a>( x: &'a [Self], batch: &'a [Self], @@ -268,16 +532,83 @@ 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 .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 +622,102 @@ 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::x86::hsum256_pd; + 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 +727,336 @@ 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(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..] + .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() + } +} + +#[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::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). + #[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(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() + } + + /// 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(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 +1166,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 +1307,447 @@ 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)); + } + } + + /// 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/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..3411d7e8dce 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -14,12 +14,16 @@ 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; +#[cfg(all( + target_arch = "x86_64", + not(all(target_feature = "avx2", target_feature = "fma")) +))] +use crate::distance::BatchIter; /// Default implementation of dot product. /// @@ -111,6 +115,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")] @@ -161,6 +183,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 +239,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 +250,128 @@ 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) + } + + 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. + // 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", + 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. + 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", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + { + 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)) + } + } +} + +/// 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(all(target_feature = "avx2", target_feature = "fma")) +))] +#[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 +/// (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 +379,246 @@ 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::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. + /// + /// 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], + 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. + #[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], + 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. + #[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 + .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 { + 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 + } + + /// 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}; @@ -285,7 +668,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( @@ -309,10 +692,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(), @@ -480,5 +862,264 @@ 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)); + } + } + + /// `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(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; + 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}" + ); + } + } + } + + // 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") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx_fma); + } + + #[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") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx); + } + + #[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") { + return; + } + check_dot_batch_kernel(x86::dot_batch_f32_avx512); } } 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..7ee2f4bc537 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -20,18 +20,41 @@ 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")] +// 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}; +#[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. /// 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)) } } @@ -52,7 +75,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) }; @@ -191,6 +213,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 +269,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,12 +280,114 @@ 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) + } + + 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. + // See `Dot::dot_batch` for f32. + #[cfg(all( + target_arch = "x86_64", + target_feature = "avx2", + 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. + 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", + not(all(target_feature = "avx2", target_feature = "fma")) + ))] + { + 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)) + } + } +} + +/// Sub-AVX2 builds: pick a `#[target_feature]` kernel once for the batch. +#[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], + 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 +/// (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 { #[inline] fn l2(x: &[Self], y: &[Self]) -> f32 { @@ -265,9 +395,277 @@ 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::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. + /// + /// 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], + 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. + #[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], + 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. + #[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 + .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 { + 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 + } + + /// 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 +1069,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] @@ -787,4 +1315,100 @@ 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(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; + 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}" + ); + } + } + } + + // 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") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx_fma); + } + + #[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") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx); + } + + #[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") { + return; + } + check_l2_batch_kernel(x86::l2_batch_f32_avx512); + } } 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..8d126d6ed57 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,166 @@ 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::x86::hsum256_ps; + 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 + } + + /// 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 +480,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.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/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..4ce7f64706d 100644 --- a/rust/lance-linalg/src/simd/f32.rs +++ b/rust/lance-linalg/src/simd/f32.rs @@ -46,14 +46,42 @@ 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. + /// + /// # Panics + /// + /// If any index is negative or lands outside `slice`. #[inline] pub fn gather(slice: &[f32], indices: &[i32; 8]) -> Self { - #[cfg(target_arch = "x86_64")] - unsafe { - use super::i32::i32x8; + // 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() + ); + } - let idx = i32x8::from(indices); - Self(_mm256_i32gather_ps::<4>(slice.as_ptr(), idx.0)) + #[cfg(target_arch = "x86_64")] + { + if is_x86_feature_detected!("avx2") { + unsafe { gather_avx2(slice, indices) } + } else { + gather_scalar_x86(slice, indices) + } } #[cfg(target_arch = "aarch64")] @@ -94,6 +122,30 @@ 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. +/// +/// 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 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 { fn from(value: &[f32]) -> Self { unsafe { Self::load_unaligned(value.as_ptr()) } @@ -439,15 +491,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 +541,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 +565,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 +581,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 +600,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 +619,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 +637,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 +653,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 +685,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 +719,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,21 +738,15 @@ 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. + 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); } } @@ -774,11 +788,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 +813,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 +836,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 +861,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 +886,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 +909,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); @@ -943,9 +933,17 @@ impl SubAssign for f32x16 { mod tests { use super::*; + use rstest::rstest; #[test] fn test_basic_ops() { + // 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!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } let a = (0..8).map(|f| f as f32).collect::>(); let b = (10..18).map(|f| f as f32).collect::>(); @@ -983,6 +981,11 @@ 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!("avx") { + 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 +1010,11 @@ 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!("avx") || !std::is_x86_feature_detected!("fma") { + return; + } let a = (0..16).map(|f| f as f32).collect::>(); let b = (10..26).map(|f| f as f32).collect::>(); @@ -1041,6 +1049,11 @@ 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!("avx") { + 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,9 +1087,59 @@ mod tests { #[test] fn test_f32x8_gather() { + // `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!("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 = 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); + } + + /// 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); + } + + /// `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); + } } diff --git a/rust/lance-linalg/src/simd/f64.rs b/rust/lance-linalg/src/simd/f64.rs index 32c0d389e5b..1276f54da56 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 / 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!("avx") { + 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,11 @@ 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!("avx") || !std::is_x86_feature_detected!("fma") { + 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 +780,11 @@ mod tests { #[test] fn test_f64x4_min() { + // `min` / `reduce_min` are AVX intrinsics. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx") { + 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 +797,11 @@ 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!("avx") { + 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 +820,11 @@ 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!("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]; let b: [f64; 8] = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]; @@ -869,6 +838,11 @@ mod tests { #[test] fn test_f64x8_min() { + // `min` / `reduce_min` are AVX intrinsics. + #[cfg(target_arch = "x86_64")] + 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]; 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(); diff --git a/rust/lance-linalg/src/simd/x86.rs b/rust/lance-linalg/src/simd/x86.rs new file mode 100644 index 00000000000..437018669e2 --- /dev/null +++ b/rust/lance-linalg/src/simd/x86.rs @@ -0,0 +1,89 @@ +// 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`. +//! +//! 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::*; + +/// 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)); + // 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) +} + +/// 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::*; + use rstest::rstest; + + #[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 sum = unsafe { hsum256_ps(_mm256_loadu_ps(lanes.as_ptr())) }; + assert_eq!(sum, expected); + } + + #[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 sum = unsafe { hsum256_pd(_mm256_loadu_pd(lanes.as_ptr())) }; + assert_eq!(sum, expected); + } +}