diff --git a/.claude/knowledge/float-accumulation-order.md b/.claude/knowledge/float-accumulation-order.md new file mode 100644 index 00000000..363f07cf --- /dev/null +++ b/.claude/knowledge/float-accumulation-order.md @@ -0,0 +1,55 @@ +# Float accumulation order across CPU tiers (2026-10-30) + +READ BY: anyone setting a float tolerance against an outside oracle +(ONNX Runtime, numpy, PyTorch), or changing `gemm_f32`, `sum`, `mean_axis`. + +Asked by lance-graph D-TRACT-HARVEST-0 (wishlist item 4): which order do +the f32 reductions accumulate in, and does it differ between backends? + +## Measured + +`examples/float_accum_probe.rs` hashes the raw bits of a fixed `gemm_f32` +(64×1024 · 1024×48) and a fixed `mean_axis` / `sum` (64×1024). Run on +five tiers; matrixmultiply picks its kernel at runtime from CPUID, so qemu +`-cpu` selects it: + +| tier (how run) | gemm kernel | `gemm_f32` | `mean_axis` | `sum` | +|---|---|---|---|---| +| host, AVX-512 | avx512 | `5bacd0823f5dcadd` | `0a7a5b9e602c690b` | `c11253de` | +| qemu Haswell | fma+avx2 | `5bacd0823f5dcadd` | `0a7a5b9e602c690b` | (not printed) | +| qemu SandyBridge | avx, no FMA | `632fcf6db284df8e` | `0a7a5b9e602c690b` | `c11253de` | +| qemu Nehalem | sse2 fallback | `632fcf6db284df8e` | `0a7a5b9e602c690b` | `c11253de` | +| qemu aarch64 | neon | `5bacd0823f5dcadd` | `0a7a5b9e602c690b` | `c11253de` | + +## What it means + +- **`sum`, `sum_axis`, `mean_axis` are bit-identical on every tier.** Over a + contiguous lane they use `numeric_util::unrolled_fold`: eight interleaved + partial sums, combined `(p0+p4)+(p1+p5)+(p2+p6)+(p3+p7)`, then the + remainder in order. The order is written in Rust, there is no FMA, and + LLVM may not reassociate IEEE adds. A non-contiguous axis sums sub-views + in axis order (`res = res + &subview`), also fixed. `mean_axis` divides + that sum by the length. +- **`gemm_f32` has two classes, not five.** The FMA kernels (AVX-512, + FMA+AVX2, NEON) agree with each other bit for bit across x86 and ARM; the + non-FMA kernels (AVX, SSE2) agree with each other. So the accumulation + order over `k` is the same in every kernel; only fused versus separate + multiply-add rounding differs. Within one tier the result is + deterministic (`matrixmultiply-threading` splits over `m`/`n`, not `k`). +- **The difference from ONNX Runtime is order, not backend.** A bound such + as lance-graph's "≤ 4 ULP on RMSNorm, all from the mean" is the gap + between the eight-way order above and ORT's own order. It does not move + between ndarray's tiers, so one tolerance holds on all of them. Matching + ORT exactly would need ORT's order (synergy-map V13), measured first. +- `masked_sum_*` are integer only and exact. + +## Pinned + +- `numeric_util::accumulation_order_tests::sum_and_mean_axis_follow_the_documented_eight_way_order`: + `sum` and `mean_axis` equal a hand-written eight-way reference, bit for + bit, and the fixture is one where a left-to-right sum differs. +- `backend::native::tests::gemm_f32_bits_fall_in_the_fma_or_non_fma_class`: + the hash must be the FMA-class value when the host takes an FMA kernel, + else the non-FMA value. A matrixmultiply upgrade that moves either class + fails here; re-measure with the probe and update both this table and the + constants. diff --git a/Cargo.toml b/Cargo.toml index 703ce1dd..20a1cb07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -54,6 +54,10 @@ required-features = ["std"] name = "vstanh_exhaustive" required-features = ["std"] +[[example]] +name = "float_accum_probe" +required-features = ["std"] + [[example]] name = "codec8_asm_probe" required-features = ["std"] diff --git a/examples/float_accum_probe.rs b/examples/float_accum_probe.rs new file mode 100644 index 00000000..15009cc2 --- /dev/null +++ b/examples/float_accum_probe.rs @@ -0,0 +1,67 @@ +//! Float accumulation order across CPU tiers: prints an FNV-1a hash of the +//! raw bits of (a) `gemm_f32` and (b) `mean_axis` / `sum` on fixed inputs, +//! plus which matrixmultiply kernel the host can take. Run the same binary +//! (or the same source) on several tiers and compare the hashes: equal +//! hashes mean bit-identical results. +//! +//! ```sh +//! cargo run --release --example float_accum_probe +//! ``` +use ndarray::{Array2, Axis}; + +fn fnv(bits: impl Iterator) -> u64 { + let mut h: u64 = 0xcbf2_9ce4_8422_2325; + for b in bits { + for byte in b.to_le_bytes() { + h ^= u64::from(byte); + h = h.wrapping_mul(0x0100_0000_01b3); + } + } + h +} + +fn main() { + // Values with full mantissas, so any change of order or of fused + // multiply-add shows up in the low bits. + let f = |i: usize, j: usize| { + (((i as u64 * 2_654_435_761 + j as u64 * 40_503) % 100_003) as f32 / 100_003.0 - 0.5) * 3.7 + }; + let (m, k, n) = (64, 1024, 48); + let a = Array2::from_shape_fn((m, k), |(i, j)| f(i, j)); + let b = Array2::from_shape_fn((k, n), |(i, j)| f(j + 7, i)); + let mut c = vec![0.0f32; m * n]; + ndarray::backend::native::gemm_f32( + m, + n, + k, + 1.0, + a.as_slice().unwrap(), + k, + b.as_slice().unwrap(), + n, + 0.0, + &mut c, + n, + ); + let mean = a.mean_axis(Axis(1)).unwrap(); + let sum: f32 = a.sum(); + let arch = if cfg!(target_arch = "aarch64") { + "aarch64" + } else { + "x86_64" + }; + #[cfg(target_arch = "x86_64")] + let tier = format!( + "avx512f={} fma={} avx2={} avx={}", + is_x86_feature_detected!("avx512f"), + is_x86_feature_detected!("fma"), + is_x86_feature_detected!("avx2"), + is_x86_feature_detected!("avx") + ); + #[cfg(not(target_arch = "x86_64"))] + let tier = String::from("neon"); + println!("arch={arch} {tier}"); + println!("gemm_f32 [{m}x{k}]·[{k}x{n}] hash={:016x}", fnv(c.iter().map(|v| v.to_bits()))); + println!("mean_axis [{m}x{k}] axis 1 hash={:016x}", fnv(mean.iter().map(|v| v.to_bits()))); + println!("sum [{m}x{k}] bits={:08x}", sum.to_bits()); +} diff --git a/src/backend/native.rs b/src/backend/native.rs index 7a5260d8..4e15b6f3 100644 --- a/src/backend/native.rs +++ b/src/backend/native.rs @@ -935,6 +935,44 @@ mod avx2 { mod tests { use super::*; + /// `gemm_f32` is matrixmultiply's runtime-selected kernel. Measured with + /// `examples/float_accum_probe.rs` on this fixture: the FMA kernels + /// (AVX-512, FMA+AVX2, NEON) give one set of bits, the non-FMA kernels + /// (AVX, SSE2 fallback) another. The order over `k` is the same in all + /// five; only fused versus separate rounding differs. Pinned so a + /// matrixmultiply change to either class fails here. + #[test] + fn gemm_f32_bits_fall_in_the_fma_or_non_fma_class() { + const FMA_CLASS: u64 = 0x5bac_d082_3f5d_cadd; + const NON_FMA_CLASS: u64 = 0x632f_cf6d_b284_df8e; + let f = |i: usize, j: usize| { + (((i as u64 * 2_654_435_761 + j as u64 * 40_503) % 100_003) as f32 / 100_003.0 - 0.5) * 3.7 + }; + let (m, k, n) = (64, 1024, 48); + let a: Vec = (0..m * k).map(|x| f(x / k, x % k)).collect(); + let b: Vec = (0..k * n).map(|x| f(x % n + 7, x / n)).collect(); + let mut c = vec![0.0f32; m * n]; + gemm_f32(m, n, k, 1.0, &a, k, &b, n, 0.0, &mut c, n); + let mut h: u64 = 0xcbf2_9ce4_8422_2325; + for v in &c { + for byte in v.to_bits().to_le_bytes() { + h ^= u64::from(byte); + h = h.wrapping_mul(0x0100_0000_01b3); + } + } + #[cfg(target_arch = "x86_64")] + let fma = std::is_x86_feature_detected!("fma") && std::is_x86_feature_detected!("avx2"); + #[cfg(target_arch = "aarch64")] + let fma = true; + #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] + return; + #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] + { + let want = if fma { FMA_CLASS } else { NON_FMA_CLASS }; + assert_eq!(h, want, "gemm_f32 bits moved (fma kernel: {fma}); re-measure with float_accum_probe"); + } + } + #[test] fn test_dot_f32() { let x = vec![1.0f32, 2.0, 3.0, 4.0]; diff --git a/src/hpc/amx_matmul.rs b/src/hpc/amx_matmul.rs index 65abadfd..460204b2 100644 --- a/src/hpc/amx_matmul.rs +++ b/src/hpc/amx_matmul.rs @@ -822,17 +822,78 @@ pub fn matmul_f32_bf16_fast( /// accurate than the CPU f32 GEMM at every size tried (see the table in /// [`matmul_f32_amx_split`]). /// -/// Delegates to `backend::native::gemm_f32` (matrixmultiply), which is -/// f32 end to end: no BF16 downcast, no tile emulation. +/// Calls `matrixmultiply::sgemm` (the kernel behind +/// `backend::native::gemm_f32`) with each view's own row and column strides, +/// so it reads `lhs` and `rhs` in place, in any layout (row-major, +/// column-major, a transposed, stepped or reversed view), and writes the +/// row-major `out` directly: no packing copies and no result buffer. f32 end to end: no BF16 downcast, no tile +/// emulation. The kernel packs its own blocks, so the result is the same, +/// bit for bit, whatever the input layout. pub fn matmul_f32( lhs: ArrayView2<'_, f32>, rhs: ArrayView2<'_, f32>, mut out: ArrayViewMut2<'_, f32>, ) -> Result<(), MatmulError> { let (m, n, k) = check_shapes(&lhs, &rhs, &out)?; - let a_f32 = pack_contig(&lhs); - let b_f32 = pack_contig(&rhs); - let mut c = vec![0.0f32; m * n]; - crate::backend::native::gemm_f32(m, n, k, 1.0, &a_f32, k, &b_f32, n, 0.0, &mut c, n); - write_contig(&mut out, &c); + if m == 0 || n == 0 { + return Ok(()); + } + let (rsa, csa) = (lhs.strides()[0], lhs.strides()[1]); + let (rsb, csb) = (rhs.strides()[0], rhs.strides()[1]); + let (rsc, csc) = (out.strides()[0], out.strides()[1]); + // SAFETY: the pointers and strides come from live views whose shapes + // `check_shapes` matched to (m, k), (k, n) and (m, n); matrixmultiply + // reads only those elements. `out` is a `&mut` view, so its elements do + // not alias each other or the inputs. beta = 0 overwrites `out` without + // reading it. + unsafe { + matrixmultiply::sgemm( + m, + k, + n, + 1.0, + lhs.as_ptr(), + rsa, + csa, + rhs.as_ptr(), + rsb, + csb, + 0.0, + out.as_mut_ptr(), + rsc, + csc, + ); + } + Ok(()) +} + +/// [`matmul_f32`] over a leading batch axis: `out[i] = lhs[i] · rhs[i]`, +/// shapes `[b, m, k] × [b, k, n] → [b, m, n]`. A `rhs` with batch size 1 is +/// shared by every `lhs[i]` (a weight matrix applied to every head). Each +/// batch entry goes through [`matmul_f32`], so it is copy-free too. +/// +/// # Errors +/// +/// [`MatmulError::ShapeMismatch`] if the batch sizes disagree (other than a +/// shared `rhs`) or any batch entry's shapes do not multiply. +pub fn matmul_f32_batched( + lhs: crate::ArrayView3<'_, f32>, rhs: crate::ArrayView3<'_, f32>, mut out: crate::ArrayViewMut3<'_, f32>, +) -> Result<(), MatmulError> { + use crate::Axis; + let (b, m, k) = lhs.dim(); + let (rb, kr, n) = rhs.dim(); + let (ob, mo, no) = out.dim(); + // Inner shapes are checked here, not only per entry: an empty batch + // runs no entry, and must not let mismatched shapes pass as `Ok`. + if ob != b || !(rb == b || rb == 1) || k != kr || m != mo || n != no { + return Err(MatmulError::ShapeMismatch { + lhs: (m, k), + rhs: (kr, n), + out: (mo, no), + }); + } + for i in 0..b { + let r = rhs.index_axis(Axis(0), if rb == 1 { 0 } else { i }); + matmul_f32(lhs.index_axis(Axis(0), i), r, out.index_axis_mut(Axis(0), i))?; + } Ok(()) } @@ -1034,6 +1095,84 @@ fn subtract_i8_to_u8_bias(c: &mut [i32], b_i8: &[i8], m: usize, n: usize, k: usi mod tests { use super::*; + /// The old `matmul_f32`: pack both inputs row-major, gemm into a buffer, + /// copy into `out`. The strided path must give the same bits. + fn matmul_f32_packed(lhs: ArrayView2<'_, f32>, rhs: ArrayView2<'_, f32>) -> crate::Array2 { + let (m, k) = lhs.dim(); + let n = rhs.dim().1; + let (a, b) = (pack_contig(&lhs), pack_contig(&rhs)); + let mut c = vec![0.0f32; m * n]; + crate::backend::native::gemm_f32(m, n, k, 1.0, &a, k, &b, n, 0.0, &mut c, n); + crate::Array2::from_shape_vec((m, n), c).unwrap() + } + + #[test] + fn matmul_f32_reads_any_layout_in_place_and_matches_the_packed_path() { + use crate::{s, Array2, ShapeBuilder}; + let f = |i: usize, j: usize| ((i * 31 + j * 17) % 23) as f32 * 0.137 - 1.3; + let (m, k, n) = (37, 29, 41); + let a = Array2::from_shape_fn((m, k), |(i, j)| f(i, j)); + let b = Array2::from_shape_fn((k, n), |(i, j)| f(j + 3, i)); + let want = matmul_f32_packed(a.view(), b.view()); + // Row-major, column-major, a transposed view, a stepped slice and a + // reversed axis: five layouts, one answer. + let a_f = Array2::from_shape_fn((m, k).f(), |(i, j)| f(i, j)); + let b_t_owned = Array2::from_shape_fn((n, k), |(j, i)| f(j + 3, i)); + let a_wide = Array2::from_shape_fn((m, 2 * k), |(i, j)| if j % 2 == 0 { f(i, j / 2) } else { f32::NAN }); + let mut a_rev = Array2::from_shape_fn((m, k), |(i, j)| f(i, k - 1 - j)); + a_rev.invert_axis(crate::Axis(1)); + let cases: [(ArrayView2<'_, f32>, ArrayView2<'_, f32>); 4] = [ + (a.view(), b.view()), + (a_f.view(), b_t_owned.t()), + (a_wide.slice(s![.., ..;2]), b.view()), + (a_rev.view(), b.view()), + ]; + for (idx, (l, r)) in cases.iter().enumerate() { + // `out` starts as NaN: beta = 0 must overwrite, never read. + let mut out = Array2::from_elem((m, n), f32::NAN); + matmul_f32(l.view(), r.view(), out.view_mut()).unwrap(); + for ((i, j), &v) in out.indexed_iter() { + assert_eq!(v.to_bits(), want[[i, j]].to_bits(), "case {idx} at ({i}, {j})"); + } + } + // `out` must stay row-major (the existing `check_shapes` contract); + // inputs may have any layout. + let mut out_f = Array2::from_elem((m, n).f(), f32::NAN); + assert!(matches!(matmul_f32(a.view(), b.view(), out_f.view_mut()), Err(MatmulError::NonContiguousOutput))); + // k = 0: the product is all zeros, still without reading `out`. + let mut z = Array2::from_elem((3, 4), f32::NAN); + matmul_f32(Array2::::zeros((3, 0)).view(), Array2::::zeros((0, 4)).view(), z.view_mut()).unwrap(); + assert!(z.iter().all(|&v| v == 0.0), "{z:?}"); + } + + #[test] + fn matmul_f32_batched_matches_per_entry_and_shares_a_single_rhs() { + use crate::Array3; + let (bsz, m, k, n) = (3, 5, 7, 4); + let l = Array3::from_shape_fn((bsz, m, k), |(b, i, j)| (b * 7 + i * 3 + j) as f32 * 0.25 - 2.0); + let r = Array3::from_shape_fn((bsz, k, n), |(b, i, j)| (b + i * 5 + j * 2) as f32 * 0.125 - 1.0); + let mut out = Array3::::zeros((bsz, m, n)); + matmul_f32_batched(l.view(), r.view(), out.view_mut()).unwrap(); + let shared = r.slice(crate::s![..1, .., ..]); + let mut out_shared = Array3::::zeros((bsz, m, n)); + matmul_f32_batched(l.view(), shared, out_shared.view_mut()).unwrap(); + for b in 0..bsz { + let want = matmul_f32_packed(l.index_axis(crate::Axis(0), b), r.index_axis(crate::Axis(0), b)); + assert_eq!(out.index_axis(crate::Axis(0), b), want, "batch {b}"); + let want0 = matmul_f32_packed(l.index_axis(crate::Axis(0), b), r.index_axis(crate::Axis(0), 0)); + assert_eq!(out_shared.index_axis(crate::Axis(0), b), want0, "shared rhs, batch {b}"); + } + // A batch of 2 against 3 is refused, not truncated. + let mut bad = Array3::::zeros((bsz, m, n)); + assert!(matmul_f32_batched(l.view(), r.slice(crate::s![..2, .., ..]), bad.view_mut()).is_err()); + // An empty batch still validates the inner shapes. + let (l0, r1, mut o0) = + (Array3::::zeros((0, 2, 3)), Array3::::zeros((1, 4, 5)), Array3::::zeros((0, 9, 8))); + assert!(matmul_f32_batched(l0.view(), r1.view(), o0.view_mut()).is_err()); + let (r_ok, mut o_ok) = (Array3::::zeros((1, 3, 5)), Array3::::zeros((0, 2, 5))); + assert!(matmul_f32_batched(l0.view(), r_ok.view(), o_ok.view_mut()).is_ok()); + } + #[test] fn test_tile_config_creation() { let cfg = TileConfig::for_dpbusd(64); diff --git a/src/hpc/gather.rs b/src/hpc/gather.rs new file mode 100644 index 00000000..e45ce33d --- /dev/null +++ b/src/hpc/gather.rs @@ -0,0 +1,164 @@ +//! Row gather from a table into a caller buffer (ONNX `Gather` on axis 0, +//! e.g. an embedding lookup). +//! +//! The table is never copied: each selected row is copied once, straight +//! into its row of `out`. Indices follow the ONNX rule: an index `i` is valid +//! when `-rows <= i < rows`, and a negative index counts from the end. An +//! index outside that range is reported as an error, never a panic and never +//! a clamp. + +use crate::{ArrayView2, ArrayViewMut2, Axis}; + +/// Why [`gather_rows_f32`] refused. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GatherError { + /// `ids[position]` is outside `[-rows, rows)`. + IndexOutOfRange { + /// Position in `ids`. + position: usize, + /// The offending index. + index: i64, + /// Number of rows in the table. + rows: usize, + }, + /// `out` is not `[ids.len(), table columns]`. + ShapeMismatch { + /// Expected `out` shape. + expected: (usize, usize), + /// Actual `out` shape. + actual: (usize, usize), + }, +} + +impl core::fmt::Display for GatherError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match *self { + GatherError::IndexOutOfRange { position, index, rows } => { + write!(f, "gather: ids[{position}] = {index} is outside [-{rows}, {rows})") + } + GatherError::ShapeMismatch { expected, actual } => { + write!(f, "gather: out is {actual:?}, expected {expected:?}") + } + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for GatherError {} + +/// `out[r] = table[ids[r]]` for every `r`, with ONNX index rules. +/// +/// Every index is checked before anything is written, so on an error `out` +/// is left unchanged. Contiguous rows are copied with `copy_from_slice`; +/// other layouts are copied element by element through the views. +/// +/// # Errors +/// +/// [`GatherError::IndexOutOfRange`] for an index outside `[-rows, rows)`, +/// [`GatherError::ShapeMismatch`] if `out` is not `[ids.len(), columns]`. +/// +/// # Example +/// +/// ``` +/// use ndarray::{array, Array2}; +/// use ndarray::hpc::gather::gather_rows_f32; +/// +/// let table = array![[0.0f32, 0.5], [1.0, 1.5], [2.0, 2.5]]; +/// let mut out = Array2::::zeros((3, 2)); +/// gather_rows_f32(table.view(), &[2, 0, -1], out.view_mut()).unwrap(); +/// assert_eq!(out, array![[2.0, 2.5], [0.0, 0.5], [2.0, 2.5]]); +/// assert!(gather_rows_f32(table.view(), &[3], out.view_mut()).is_err()); +/// ``` +pub fn gather_rows_f32( + table: ArrayView2<'_, f32>, ids: &[i64], mut out: ArrayViewMut2<'_, f32>, +) -> Result<(), GatherError> { + let (rows, cols) = table.dim(); + let expected = (ids.len(), cols); + if out.dim() != expected { + return Err(GatherError::ShapeMismatch { + expected, + actual: out.dim(), + }); + } + let row_of = |position: usize, index: i64| -> Result { + let err = GatherError::IndexOutOfRange { position, index, rows }; + let r = i128::from(index); + let n = rows as i128; + if r < -n || r >= n { + return Err(err); + } + Ok(if r < 0 { (r + n) as usize } else { r as usize }) + }; + // Validate everything first, so a bad index leaves `out` untouched. + for (p, &i) in ids.iter().enumerate() { + row_of(p, i)?; + } + for (p, &i) in ids.iter().enumerate() { + let src = table.index_axis(Axis(0), row_of(p, i)?); + let mut dst = out.index_axis_mut(Axis(0), p); + match (src.as_slice(), dst.as_slice_mut()) { + (Some(s), Some(d)) => d.copy_from_slice(s), + _ => dst.assign(&src), + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{s, Array2, ShapeBuilder}; + + #[test] + fn gathers_rows_with_onnx_negative_indices_in_any_layout() { + let table = Array2::from_shape_fn((5, 3), |(i, j)| (i * 10 + j) as f32); + let ids = [4i64, 0, -1, -5, 2]; + let want = Array2::from_shape_fn((5, 3), |(r, j)| [4, 0, 4, 0, 2][r] as f32 * 10.0 + j as f32); + let mut out = Array2::::zeros((5, 3)); + gather_rows_f32(table.view(), &ids, out.view_mut()).unwrap(); + assert_eq!(out, want); + // A transposed (non-contiguous rows) table takes the view path. + let tf = Array2::from_shape_fn((5, 3).f(), |(i, j)| (i * 10 + j) as f32); + let tv = tf.view(); + assert!(tv.index_axis(Axis(0), 0).as_slice().is_none(), "fixture must take the view path"); + let mut out2 = Array2::::zeros((5, 3)); + gather_rows_f32(tv, &ids, out2.view_mut()).unwrap(); + assert_eq!(out2, want); + // A stepped `out` view. + let mut big = Array2::::from_elem((10, 3), -9.0); + gather_rows_f32(table.view(), &ids, big.slice_mut(s![..;2, ..])).unwrap(); + assert_eq!(big.slice(s![..;2, ..]), want); + assert!(big.slice(s![1..;2, ..]).iter().all(|&v| v == -9.0)); + } + + #[test] + fn refuses_out_of_range_indices_and_leaves_out_untouched() { + let table = Array2::from_shape_fn((5, 2), |(i, j)| (i + j) as f32); + for (ids, bad_pos, bad) in [ + (vec![0i64, 5], 1usize, 5i64), + (vec![-6i64], 0, -6), + (vec![1, i64::MIN], 1, i64::MIN), + (vec![i64::MAX], 0, i64::MAX), + ] { + let mut out = Array2::::from_elem((ids.len(), 2), 7.0); + assert_eq!( + gather_rows_f32(table.view(), &ids, out.view_mut()), + Err(GatherError::IndexOutOfRange { + position: bad_pos, + index: bad, + rows: 5 + }) + ); + assert!(out.iter().all(|&v| v == 7.0), "out written before the error: {out:?}"); + } + let mut wrong = Array2::::zeros((2, 3)); + assert!(matches!( + gather_rows_f32(table.view(), &[0, 1], wrong.view_mut()), + Err(GatherError::ShapeMismatch { .. }) + )); + // An empty table refuses every index, including 0 and -0. + let empty = Array2::::zeros((0, 2)); + let mut o = Array2::::zeros((1, 2)); + assert!(gather_rows_f32(empty.view(), &[0], o.view_mut()).is_err()); + } +} diff --git a/src/hpc/mod.rs b/src/hpc/mod.rs index ebe3fc38..e8fee821 100644 --- a/src/hpc/mod.rs +++ b/src/hpc/mod.rs @@ -126,6 +126,8 @@ pub mod nars; #[allow(missing_docs)] pub mod blackboard; pub mod bulk; +/// Row gather from a table into a caller buffer (ONNX `Gather`, axis 0). +pub mod gather; #[allow(missing_docs)] pub mod bnn; #[allow(missing_docs)] diff --git a/src/numeric_util.rs b/src/numeric_util.rs index 1879abdb..49e4b4c2 100644 --- a/src/numeric_util.rs +++ b/src/numeric_util.rs @@ -130,3 +130,61 @@ where true } + +/// The f32 summation order behind `sum`, `sum_axis` and `mean_axis` over a +/// contiguous lane is the eight interleaved partial sums of +/// [`unrolled_fold`], combined as `(p0+p4)+(p1+p5)+(p2+p6)+(p3+p7)`, then the +/// remaining elements in order. It is written out in Rust, has no fused +/// multiply-add, and LLVM may not reassociate IEEE adds, so it is the same on +/// every CPU tier. `examples/float_accum_probe.rs` measured identical bits on +/// AVX-512, FMA+AVX2, AVX, SSE2 and NEON. +#[cfg(test)] +mod accumulation_order_tests { + use crate::{Array2, Axis}; + + fn fixture() -> Array2 { + Array2::from_shape_fn((64, 1027), |(i, j)| { + (((i as u64 * 2_654_435_761 + j as u64 * 40_503) % 100_003) as f32 / 100_003.0 - 0.5) * 3.7 + }) + } + + /// The order spelled out by hand, independent of `unrolled_fold`. + fn eight_way(xs: &[f32]) -> f32 { + let mut p = [0.0f32; 8]; + let full = xs.len() / 8 * 8; + for c in xs[..full].as_chunks::<8>().0 { + for (pk, &x) in p.iter_mut().zip(c) { + *pk += x; + } + } + let mut acc = 0.0f32; + for (a, b) in [(0, 4), (1, 5), (2, 6), (3, 7)] { + acc += p[a] + p[b]; + } + for &x in &xs[full..] { + acc += x; + } + acc + } + + #[test] + fn sum_and_mean_axis_follow_the_documented_eight_way_order() { + let a = fixture(); + // 1027 per row: 128 full groups of 8 plus a 3-element remainder. + for (r, row) in a.rows().into_iter().enumerate() { + let s = row.as_slice().unwrap(); + assert_eq!(row.sum().to_bits(), eight_way(s).to_bits(), "row {r} sum"); + } + let mean = a.mean_axis(Axis(1)).unwrap(); + for (r, row) in a.rows().into_iter().enumerate() { + let want = eight_way(row.as_slice().unwrap()) / 1027.0; + assert_eq!(mean[r].to_bits(), want.to_bits(), "row {r} mean"); + } + let all = a.as_slice().unwrap(); + assert_eq!(a.sum().to_bits(), eight_way(all).to_bits(), "whole-array sum"); + // Anti-vacuity: a plain left-to-right sum differs on this fixture, so + // the comparison above can tell the two orders apart. + let seq: f32 = all.iter().sum(); + assert_ne!(seq.to_bits(), a.sum().to_bits(), "fixture cannot separate the orders"); + } +} diff --git a/src/simd.rs b/src/simd.rs index 348961f2..da575033 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -880,12 +880,20 @@ pub use crate::simd_masking_ops::{ // 2026-09-16: the ordered u64 family (N3/G2). The width where the packing // is NOT free — `U64x8` is 8 lanes, so eight groups share a word and each // lands at its own byte. + blend_f32, + blend_f32_scalar, blend_i32, blend_i32_scalar, + blend_i64, + blend_i64_scalar, + eq_f32_pair_to_mask, eq_f32_to_mask, eq_f64_to_mask, + eq_i32_pair_to_mask, eq_i32_to_mask, eq_i32_to_mask_under, + eq_i64_pair_to_mask, + eq_i64_to_mask, eq_u32_strided_to_mask, eq_u32_to_mask, eq_u32_to_mask_under, @@ -894,28 +902,44 @@ pub use crate::simd_masking_ops::{ eq_u8_to_mask, fold_bounded_cross_power_sums_tiles, fold_bounded_power_sums_tiles, + ge_f32_pair_to_mask, ge_f32_to_mask, ge_f64_to_mask, + ge_i32_pair_to_mask, ge_i32_to_mask, ge_i32_to_mask_under, + ge_i64_pair_to_mask, + ge_i64_to_mask, ge_u64_to_mask, ge_u8_to_mask, + gt_f32_pair_to_mask, gt_f32_to_mask, gt_f64_to_mask, + gt_i32_pair_to_mask, gt_i32_to_mask, gt_i32_to_mask_under, + gt_i64_pair_to_mask, + gt_i64_to_mask, gt_u64_to_mask, gt_u8_to_mask, + le_f32_pair_to_mask, le_f32_to_mask, le_f64_to_mask, + le_i32_pair_to_mask, le_i32_to_mask, le_i32_to_mask_under, + le_i64_pair_to_mask, + le_i64_to_mask, le_u64_to_mask, le_u8_to_mask, + lt_f32_pair_to_mask, lt_f32_to_mask, lt_f64_to_mask, + lt_i32_pair_to_mask, lt_i32_to_mask, lt_i32_to_mask_under, + lt_i64_pair_to_mask, + lt_i64_to_mask, lt_u64_to_mask, lt_u8_to_mask, mask_all, @@ -980,6 +1004,7 @@ pub use crate::simd_masking_ops::{ ne_f64_to_mask, ne_i32_to_mask, ne_i32_to_mask_under, + ne_i64_to_mask, ne_u32_to_mask, ne_u32_to_mask_under, ne_u64_to_mask, diff --git a/src/simd_masking_ops.rs b/src/simd_masking_ops.rs index 1cb80850..ebfa444c 100644 --- a/src/simd_masking_ops.rs +++ b/src/simd_masking_ops.rs @@ -5799,6 +5799,251 @@ fn masked_fold_i32(values: &[i32], mask_words: &[u64], f: impl Fn(i32, i32) -> i acc } +// ──────────────────────────────────────────────────────────────────────── +// Two-array compares (ONNX `Equal`, `Less`, `Greater`, … with two tensor +// operands) and the `i64` family +// ──────────────────────────────────────────────────────────────────────── +// +// The scalar-threshold `*_to_mask` words above compare a lane against one +// value. ONNX's comparison ops take two tensors, so this family compares +// `a[i]` against `b[i]`. Same output contract: bit `i % 64` of word `i / 64`, +// full overwrite, trailing bits zero. +// +// The `f32` pair words are IEEE, not SQL: a NaN on either side compares +// false for every relation (ONNX semantics). They do not use the NULL=NaN +// plane of `lt_f32_to_mask` and its siblings. +// +// `i64` has no signed compare lane on every backend, so the `i64` words +// compare through `U64x8::cmpgt_mask` / `cmpeq_mask` after flipping the sign +// bit: `x ^ (1 << 63)` maps signed order onto unsigned order exactly +// (`i64::MIN` → 0, `-1` → `2^63 − 1`, `0` → `2^63`, `i64::MAX` → `u64::MAX`). + +/// [`pack`] over two slices of equal length. +fn pack_pair( + name: &str, a: &[T], b: &[T], out_words: &mut [u64], group_bits: impl Fn([T; L], [T; L], usize) -> u64, +) { + assert_eq!(a.len(), b.len(), "{name}: a/b length mismatch"); + const { assert!(64 % L == 0, "a mask word must hold whole lane groups") } + let k = 64 / L; + let n = a.len(); + let words = mask_words_for(n); + assert!(out_words.len() >= words, "{name}: out_words.len()={} < required {words}", out_words.len()); + for w in out_words.iter_mut() { + *w = 0; + } + let (ca, ta) = a.as_chunks::(); + let (cb, tb) = b.as_chunks::(); + for (g, (x, y)) in ca.iter().zip(cb).enumerate() { + out_words[g / k] |= group_bits(*x, *y, L) << ((g % k) * L); + } + if !ta.is_empty() { + let g = ca.len(); + let live = ta.len(); + let bits = group_bits(pad_tail(ta), pad_tail(tb), live); + debug_assert!(live == 64 || bits >> live == 0, "{name}: tail bits above live={live} must be zero"); + out_words[g / k] |= bits << ((g % k) * L); + } +} + +/// The sign-flipped `u64` image of eight `i64` lanes (see the section note). +#[inline(always)] +fn i64_ordered(lanes: [i64; 8]) -> crate::simd::U64x8 { + crate::simd::U64x8::from_array(lanes.map(|x| (x as u64) ^ (1 << 63))) +} + +macro_rules! pair_cmp_i32 { + ($(#[$doc:meta])* $name:ident, |$x:ident, $y:ident| $rel:expr) => { + $(#[$doc])* + /// + /// # Panics + /// + /// Panics unless `a.len() == b.len()` and + /// `out_words.len() >= a.len().div_ceil(64)`. + #[inline] + pub fn $name(a: &[i32], b: &[i32], out_words: &mut [u64]) { + pack_pair::(stringify!($name), a, b, out_words, |xa, ya, live| { + let $x = crate::simd::I32x16::from_array(xa); + let $y = crate::simd::I32x16::from_array(ya); + (($rel) & live16(live)) as u64 + }); + } + }; +} + +pair_cmp_i32!( + /// Packs `a[i] == b[i]` over `i32` (ONNX `Equal`). + /// + /// ``` + /// use ndarray::simd::eq_i32_pair_to_mask; + /// let mut w = [0u64; 1]; + /// eq_i32_pair_to_mask(&[1, 2, 3], &[1, 0, 3], &mut w); + /// assert_eq!(w[0], 0b101); + /// ``` + eq_i32_pair_to_mask, |x, y| !(x.gt_bitmask(y) | y.gt_bitmask(x)) +); +pair_cmp_i32!( + /// Packs `a[i] < b[i]` over `i32` (signed; ONNX `Less`). + lt_i32_pair_to_mask, |x, y| y.gt_bitmask(x) +); +pair_cmp_i32!( + /// Packs `a[i] <= b[i]` over `i32` (signed; ONNX `LessOrEqual`). + le_i32_pair_to_mask, |x, y| !x.gt_bitmask(y) +); +pair_cmp_i32!( + /// Packs `a[i] > b[i]` over `i32` (signed; ONNX `Greater`). + gt_i32_pair_to_mask, |x, y| x.gt_bitmask(y) +); +pair_cmp_i32!( + /// Packs `a[i] >= b[i]` over `i32` (signed; ONNX `GreaterOrEqual`). + ge_i32_pair_to_mask, |x, y| !y.gt_bitmask(x) +); + +macro_rules! pair_cmp_i64 { + ($(#[$doc:meta])* $name:ident, |$x:ident, $y:ident| $rel:expr) => { + $(#[$doc])* + /// + /// # Panics + /// + /// Panics unless `a.len() == b.len()` and + /// `out_words.len() >= a.len().div_ceil(64)`. + #[inline] + pub fn $name(a: &[i64], b: &[i64], out_words: &mut [u64]) { + pack_pair::(stringify!($name), a, b, out_words, |xa, ya, live| { + let $x = i64_ordered(xa); + let $y = i64_ordered(ya); + (($rel) & live8(live)) as u64 + }); + } + }; +} + +pair_cmp_i64!( + /// Packs `a[i] == b[i]` over `i64` (ONNX `Equal`). + /// + /// ``` + /// use ndarray::simd::eq_i64_pair_to_mask; + /// let mut w = [0u64; 1]; + /// eq_i64_pair_to_mask(&[i64::MIN, 2, -3], &[i64::MIN, 0, -3], &mut w); + /// assert_eq!(w[0], 0b101); + /// ``` + eq_i64_pair_to_mask, |x, y| x.cmpeq_mask(y) +); +pair_cmp_i64!( + /// Packs `a[i] < b[i]` over `i64` (signed; ONNX `Less`). + lt_i64_pair_to_mask, |x, y| y.cmpgt_mask(x) +); +pair_cmp_i64!( + /// Packs `a[i] <= b[i]` over `i64` (signed; ONNX `LessOrEqual`). + le_i64_pair_to_mask, |x, y| !x.cmpgt_mask(y) +); +pair_cmp_i64!( + /// Packs `a[i] > b[i]` over `i64` (signed; ONNX `Greater`). + gt_i64_pair_to_mask, |x, y| x.cmpgt_mask(y) +); +pair_cmp_i64!( + /// Packs `a[i] >= b[i]` over `i64` (signed; ONNX `GreaterOrEqual`). + ge_i64_pair_to_mask, |x, y| !y.cmpgt_mask(x) +); + +macro_rules! pair_cmp_f32 { + ($(#[$doc:meta])* $name:ident, $rel:ident) => { + $(#[$doc])* + /// + /// IEEE: a NaN on either side compares false. + /// + /// # Panics + /// + /// Panics unless `a.len() == b.len()` and + /// `out_words.len() >= a.len().div_ceil(64)`. + #[inline] + pub fn $name(a: &[f32], b: &[f32], out_words: &mut [u64]) { + pack_pair::(stringify!($name), a, b, out_words, |xa, ya, live| { + let x = crate::simd::F32x16::from_array(xa); + let y = crate::simd::F32x16::from_array(ya); + (x.$rel(y).to_bitmask() & live16(live)) as u64 + }); + } + }; +} + +pair_cmp_f32!( + /// Packs `a[i] == b[i]` over `f32` (ONNX `Equal`; `-0.0 == 0.0`). + /// + /// ``` + /// use ndarray::simd::eq_f32_pair_to_mask; + /// let mut w = [0u64; 1]; + /// eq_f32_pair_to_mask(&[1.0, f32::NAN, -0.0], &[1.0, f32::NAN, 0.0], &mut w); + /// assert_eq!(w[0], 0b101); + /// ``` + eq_f32_pair_to_mask, simd_eq +); +pair_cmp_f32!( + /// Packs `a[i] < b[i]` over `f32` (ONNX `Less`). + lt_f32_pair_to_mask, simd_lt +); +pair_cmp_f32!( + /// Packs `a[i] <= b[i]` over `f32` (ONNX `LessOrEqual`). + le_f32_pair_to_mask, simd_le +); +pair_cmp_f32!( + /// Packs `a[i] > b[i]` over `f32` (ONNX `Greater`). + gt_f32_pair_to_mask, simd_gt +); +pair_cmp_f32!( + /// Packs `a[i] >= b[i]` over `f32` (ONNX `GreaterOrEqual`). + ge_f32_pair_to_mask, simd_ge +); + +macro_rules! scalar_cmp_i64 { + ($(#[$doc:meta])* $name:ident, |$x:ident, $t:ident| $rel:expr) => { + $(#[$doc])* + /// + /// # Panics + /// + /// Panics if `out_words.len() < values.len().div_ceil(64)`. + #[inline] + pub fn $name(values: &[i64], threshold: i64, out_words: &mut [u64]) { + let $t = i64_ordered([threshold; 8]); + pack::(stringify!($name), values, out_words, |lanes, live| { + let $x = i64_ordered(lanes); + (($rel) & live8(live)) as u64 + }); + } + }; +} + +scalar_cmp_i64!( + /// Packs `values[i] == threshold` over `i64`. + /// + /// ``` + /// use ndarray::simd::eq_i64_to_mask; + /// let mut w = [0u64; 1]; + /// eq_i64_to_mask(&[-1, 5, -1, i64::MAX], -1, &mut w); + /// assert_eq!(w[0], 0b0101); + /// ``` + eq_i64_to_mask, |x, t| x.cmpeq_mask(t) +); +scalar_cmp_i64!( + /// Packs `values[i] != threshold` over `i64`. + ne_i64_to_mask, |x, t| !x.cmpeq_mask(t) +); +scalar_cmp_i64!( + /// Packs `values[i] < threshold` over `i64` (signed). + lt_i64_to_mask, |x, t| t.cmpgt_mask(x) +); +scalar_cmp_i64!( + /// Packs `values[i] <= threshold` over `i64` (signed). + le_i64_to_mask, |x, t| !x.cmpgt_mask(t) +); +scalar_cmp_i64!( + /// Packs `values[i] > threshold` over `i64` (signed). + gt_i64_to_mask, |x, t| x.cmpgt_mask(t) +); +scalar_cmp_i64!( + /// Packs `values[i] >= threshold` over `i64` (signed). + ge_i64_to_mask, |x, t| !t.cmpgt_mask(x) +); + /// `dst[i] = if mask bit i { a[i] } else { b[i] }` — the conditional-select /// (`CASE WHEN`) over a row mask, with no compaction. Bits at or past /// `a.len()` are ignored. Plain index loop, deliberately: a bit-per-element @@ -5825,18 +6070,34 @@ fn masked_fold_i32(values: &[i32], mask_words: &[u64], f: impl Fn(i32, i32) -> i /// ``` #[inline] pub fn blend_i32(mask_words: &[u64], a: &[i32], b: &[i32], dst: &mut [i32]) { - assert_eq!(a.len(), b.len(), "blend_i32: a/b length mismatch"); - assert_eq!(a.len(), dst.len(), "blend_i32: a/dst length mismatch"); + blend_slices("blend_i32", mask_words, a, b, dst); +} + +/// Shared body of the `blend_*` family: `dst[i] = bit i ? a[i] : b[i]`. +/// Generic over the element type because the select moves whole elements +/// and never looks at their bits: an `f32` NaN payload or `-0.0` comes +/// through unchanged. +#[inline] +fn blend_slices(name: &str, mask_words: &[u64], a: &[T], b: &[T], dst: &mut [T]) { + assert_eq!(a.len(), b.len(), "{name}: a/b length mismatch"); + blend_with(name, mask_words, a, |i| b[i], dst); +} + +/// [`blend_slices`] with the `else` side computed per index (a slice +/// element, or one broadcast scalar). +#[inline] +fn blend_with(name: &str, mask_words: &[u64], a: &[T], b: impl Fn(usize) -> T, dst: &mut [T]) { + assert_eq!(a.len(), dst.len(), "{name}: a/dst length mismatch"); let n = a.len(); assert!( mask_words.len() >= mask_words_for(n), - "blend_i32: mask_words.len()={} < required {}", + "{name}: mask_words.len()={} < required {}", mask_words.len(), mask_words_for(n) ); for i in 0..n { let bit = (mask_words[i / 64] >> (i % 64)) & 1; - dst[i] = if bit == 1 { a[i] } else { b[i] }; + dst[i] = if bit == 1 { a[i] } else { b(i) }; } } @@ -5863,18 +6124,96 @@ pub fn blend_i32(mask_words: &[u64], a: &[i32], b: &[i32], dst: &mut [i32]) { /// ``` #[inline] pub fn blend_i32_scalar(mask_words: &[u64], a: &[i32], b: i32, dst: &mut [i32]) { - assert_eq!(a.len(), dst.len(), "blend_i32_scalar: a/dst length mismatch"); - let n = a.len(); - assert!( - mask_words.len() >= mask_words_for(n), - "blend_i32_scalar: mask_words.len()={} < required {}", - mask_words.len(), - mask_words_for(n) - ); - for i in 0..n { - let bit = (mask_words[i / 64] >> (i % 64)) & 1; - dst[i] = if bit == 1 { a[i] } else { b }; - } + blend_with("blend_i32_scalar", mask_words, a, |_| b, dst); +} + +/// `dst[i] = if mask bit i { a[i] } else { b[i] }` over `f32`: ONNX +/// `Where(C, X, Y)` with float `X`, `Y`. Elements are moved, never computed +/// on, so NaN payloads and `-0.0` come through bit for bit. Same contract as +/// [`blend_i32`]. +/// +/// # Panics +/// +/// Panics unless `a.len() == b.len() == dst.len()` and +/// `mask_words.len() >= a.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::blend_f32; +/// +/// let mut dst = [0.0f32; 3]; +/// blend_f32(&[0b101], &[1.0, 2.0, 3.0], &[-1.0, -2.0, -3.0], &mut dst); +/// assert_eq!(dst, [1.0, -2.0, 3.0]); +/// ``` +#[inline] +pub fn blend_f32(mask_words: &[u64], a: &[f32], b: &[f32], dst: &mut [f32]) { + blend_slices("blend_f32", mask_words, a, b, dst); +} + +/// [`blend_f32`] with a scalar `else` value: the attention-mask shape +/// `Where(mask, scores, -inf)` without filling a `-inf` vector. +/// +/// # Panics +/// +/// Panics unless `a.len() == dst.len()` and +/// `mask_words.len() >= a.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::blend_f32_scalar; +/// +/// let mut dst = [0.0f32; 3]; +/// blend_f32_scalar(&[0b011], &[0.5, 0.25, 0.125], f32::NEG_INFINITY, &mut dst); +/// assert_eq!(dst, [0.5, 0.25, f32::NEG_INFINITY]); +/// ``` +#[inline] +pub fn blend_f32_scalar(mask_words: &[u64], a: &[f32], b: f32, dst: &mut [f32]) { + blend_with("blend_f32_scalar", mask_words, a, |_| b, dst); +} + +/// `dst[i] = if mask bit i { a[i] } else { b[i] }` over `i64`: ONNX +/// `Where` on index tensors. Same contract as [`blend_i32`]. +/// +/// # Panics +/// +/// Panics unless `a.len() == b.len() == dst.len()` and +/// `mask_words.len() >= a.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::blend_i64; +/// +/// let mut dst = [0i64; 2]; +/// blend_i64(&[0b10], &[1, i64::MAX], &[i64::MIN, 2], &mut dst); +/// assert_eq!(dst, [i64::MIN, i64::MAX]); +/// ``` +#[inline] +pub fn blend_i64(mask_words: &[u64], a: &[i64], b: &[i64], dst: &mut [i64]) { + blend_slices("blend_i64", mask_words, a, b, dst); +} + +/// [`blend_i64`] with a scalar `else` value. +/// +/// # Panics +/// +/// Panics unless `a.len() == dst.len()` and +/// `mask_words.len() >= a.len().div_ceil(64)`. +/// +/// # Examples +/// +/// ``` +/// use ndarray::simd::blend_i64_scalar; +/// +/// let mut dst = [0i64; 2]; +/// blend_i64_scalar(&[0b01], &[7, 8], -1, &mut dst); +/// assert_eq!(dst, [7, -1]); +/// ``` +#[inline] +pub fn blend_i64_scalar(mask_words: &[u64], a: &[i64], b: i64, dst: &mut [i64]) { + blend_with("blend_i64_scalar", mask_words, a, |_| b, dst); } // ──────────────────────────────────────────────────────────────────────── @@ -8814,6 +9153,146 @@ mod tests { assert_eq!(masked_max_i32(&v, &[u64::MAX]), Some(9)); } + /// Every pair compare and every `i64` scalar compare against a plain + /// scalar loop, over lengths that end mid-group and mid-word, with the + /// extremes that a wrong sign handling or a wrong NaN rule gets wrong. + #[test] + fn pair_and_i64_compares_match_scalar_loops() { + fn want(a: &[T], b: &[T], rel: impl Fn(T, T) -> bool) -> Vec { + let mut w = vec![0u64; a.len().div_ceil(64)]; + for i in 0..a.len() { + if rel(a[i], b[i]) { + w[i / 64] |= 1 << (i % 64); + } + } + w + } + type F = fn(&[T], &[T], &mut [u64]); + type Case = (F, fn(T, T) -> bool); + type S = fn(&[i64], i64, &mut [u64]); + type ScalarCase = (S, fn(i64, i64) -> bool); + for n in [0usize, 1, 7, 8, 15, 16, 63, 64, 65, 131] { + let ext_i32 = [i32::MIN, -1, 0, 1, i32::MAX]; + let ext_i64 = [i64::MIN, i64::MIN + 1, -1, 0, 1, i64::MAX - 1, i64::MAX]; + let ext_f32 = [f32::NEG_INFINITY, -1.0, -0.0, 0.0, 1.0, f32::INFINITY, f32::NAN]; + let ai: Vec = (0..n).map(|i| ext_i32[(i * 3) % 5]).collect(); + let bi: Vec = (0..n).map(|i| ext_i32[(i * 7 + 1) % 5]).collect(); + let al: Vec = (0..n).map(|i| ext_i64[(i * 3) % 7]).collect(); + let bl: Vec = (0..n).map(|i| ext_i64[(i * 5 + 2) % 7]).collect(); + let af: Vec = (0..n).map(|i| ext_f32[(i * 3) % 7]).collect(); + let bf: Vec = (0..n).map(|i| ext_f32[(i * 5 + 1) % 7]).collect(); + let words = n.div_ceil(64); + let run = |f: &dyn Fn(&mut [u64])| { + let mut w = vec![!0u64; words]; + f(&mut w); + w + }; + let i32s: [Case; 5] = [ + (eq_i32_pair_to_mask, |x, y| x == y), + (lt_i32_pair_to_mask, |x, y| x < y), + (le_i32_pair_to_mask, |x, y| x <= y), + (gt_i32_pair_to_mask, |x, y| x > y), + (ge_i32_pair_to_mask, |x, y| x >= y), + ]; + for (k, (f, r)) in i32s.iter().enumerate() { + assert_eq!(run(&|w| f(&ai, &bi, w)), want(&ai, &bi, r), "i32 pair #{k}, n={n}"); + } + let i64s: [Case; 5] = [ + (eq_i64_pair_to_mask, |x, y| x == y), + (lt_i64_pair_to_mask, |x, y| x < y), + (le_i64_pair_to_mask, |x, y| x <= y), + (gt_i64_pair_to_mask, |x, y| x > y), + (ge_i64_pair_to_mask, |x, y| x >= y), + ]; + for (k, (f, r)) in i64s.iter().enumerate() { + assert_eq!(run(&|w| f(&al, &bl, w)), want(&al, &bl, r), "i64 pair #{k}, n={n}"); + } + let f32s: [Case; 5] = [ + (eq_f32_pair_to_mask, |x, y| x == y), + (lt_f32_pair_to_mask, |x, y| x < y), + (le_f32_pair_to_mask, |x, y| x <= y), + (gt_f32_pair_to_mask, |x, y| x > y), + (ge_f32_pair_to_mask, |x, y| x >= y), + ]; + for (k, (f, r)) in f32s.iter().enumerate() { + assert_eq!(run(&|w| f(&af, &bf, w)), want(&af, &bf, r), "f32 pair #{k}, n={n}"); + } + let scalars: [ScalarCase; 6] = [ + (eq_i64_to_mask, |x, t| x == t), + (ne_i64_to_mask, |x, t| x != t), + (lt_i64_to_mask, |x, t| x < t), + (le_i64_to_mask, |x, t| x <= t), + (gt_i64_to_mask, |x, t| x > t), + (ge_i64_to_mask, |x, t| x >= t), + ]; + for t in ext_i64 { + let tv = vec![t; n]; + for (k, (f, r)) in scalars.iter().enumerate() { + assert_eq!(run(&|w| f(&al, t, w)), want(&al, &tv, r), "i64 scalar #{k}, t={t}, n={n}"); + } + } + } + } + + /// The mask for the blend tests: 131 lanes (two full words, a partial + /// third), stray bits past the end, both sides taken. + fn blend_test_mask(n: usize) -> [u64; 3] { + let mut m = [0u64; 3]; + for i in 0..n { + if (i * 37) % 5 < 2 { + m[i / 64] |= 1 << (i % 64); + } + } + m[2] |= !0u64 << (n % 64); + m + } + + #[test] + fn blend_f32_and_i64_select_by_bit_and_move_bits_unchanged() { + let n = 131; + let m = blend_test_mask(n); + let bit = |i: usize| (m[i / 64] >> (i % 64)) & 1 == 1; + // f32: a NaN with a payload and -0.0 on both sides must arrive bit + // for bit; a select that computed on its operands would canonicalise + // the NaN or turn -0.0 into +0.0. + let nan = f32::from_bits(0x7fc0_1234); + let a: Vec = (0..n) + .map(|i| if i % 7 == 0 { nan } else { i as f32 * 0.5 }) + .collect(); + let b: Vec = (0..n) + .map(|i| if i % 3 == 0 { -0.0 } else { -(i as f32) }) + .collect(); + let mut d = vec![0.0f32; n]; + blend_f32(&m, &a, &b, &mut d); + for i in 0..n { + let want = if bit(i) { a[i] } else { b[i] }; + assert_eq!(d[i].to_bits(), want.to_bits(), "f32 lane {i}"); + } + let mut ds = vec![0.0f32; n]; + blend_f32_scalar(&m, &a, f32::NEG_INFINITY, &mut ds); + for i in 0..n { + let want = if bit(i) { a[i] } else { f32::NEG_INFINITY }; + assert_eq!(ds[i].to_bits(), want.to_bits(), "f32 scalar lane {i}"); + } + // i64: the extremes, which an i32 path would truncate. + let a: Vec = (0..n as i64) + .map(|i| if i % 5 == 0 { i64::MAX } else { i << 40 }) + .collect(); + let b: Vec = (0..n as i64) + .map(|i| if i % 4 == 0 { i64::MIN } else { -i }) + .collect(); + let mut d = vec![0i64; n]; + blend_i64(&m, &a, &b, &mut d); + let mut ds = vec![0i64; n]; + blend_i64_scalar(&m, &a, i64::MIN + 1, &mut ds); + for i in 0..n { + assert_eq!(d[i], if bit(i) { a[i] } else { b[i] }, "i64 lane {i}"); + assert_eq!(ds[i], if bit(i) { a[i] } else { i64::MIN + 1 }, "i64 scalar lane {i}"); + } + let taken = (0..n).filter(|&i| bit(i)).count(); + assert!(taken > 0 && taken < n, "fixture must take both sides: {taken}"); + } + #[test] fn blend_i32_scalar_equals_blend_i32_with_a_filled_vector() { // 131 lanes: two full words and a partial third, with stray bits set diff --git a/src/simd_stride.rs b/src/simd_stride.rs index ecf7b15a..210b65ec 100644 --- a/src/simd_stride.rs +++ b/src/simd_stride.rs @@ -6,8 +6,9 @@ //! assembled from shuffles and masks (CLAUDE.md, FACET STRIDE). //! //! This module holds the reusable forms of that shape: [`lut_u8`] for a -//! plain byte table and [`lut9`] for an odd function over a sign-magnitude -//! 9-bit (SM9) lane, u8 magnitude plus sign plane, so callers do not +//! plain byte table and [`lut9`] for an odd function over a canonical DR512 +//! lane (u8 magnitude plus ONE presence/polarity plane: P=0,M=0 is "never +//! touched", not a second sign), so callers do not //! rebuild them. Like [`crate::simd_masking_ops`] it names no ISA and owns //! no backend; it is slice ergonomics over the position view.