Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions src/simd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -683,6 +683,25 @@ pub use crate::hpc::bf16_tile_gemm::{
#[cfg(target_arch = "x86_64")]
pub use crate::simd_amx::{amx_report, cpu_model, CpuModel};

// Packed-bitmask predicates + mask algebra — the columnar-selection lane.
// Slice-level siblings of `add_i8` / `dot_i8`, built on the lane-level
// `U32x16::eq_bitmask` / `I32x16::gt_bitmask` methods. Surfaced here because
// the W1a invariant is "all SIMD from `ndarray::simd`": a consumer that had to
// reach into `ndarray::simd_int_ops` (or worse, write its own compare-and-pack
// loop) would be a polyfill bypass. Bit order is normative and identical
// across all of them — element `i` at bit `i % 64` of word `i / 64`, trailing
// bits zero. See `src/simd_int_ops.rs` for the full statement.
#[cfg(feature = "std")]
pub use crate::simd_int_ops::{
eq_u32_strided_to_mask, eq_u32_to_mask, gt_i32_to_mask, mask_and, mask_and_assign, mask_or, mask_or_assign,
masked_sum_i32,
};
// The popcount that closes the loop on the masks above: `mask_count` in ABI
// terms. Already public at `ndarray::bitwise::popcount_batch_u64`; re-exported
// here so a mask producer and its reducer share one import path (the sibling
// `popcount_raw` / `hamming_distance_raw` re-export is just above).
pub use crate::bitwise::popcount_batch_u64;

// Elementwise slice ops — polyfill-dispatched (F32x16/F64x8 chunks + scalar tail).
#[cfg(feature = "std")]
pub use crate::simd_ops::{
Expand Down
51 changes: 51 additions & 0 deletions src/simd_avx2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1764,6 +1764,29 @@ impl U32x16 {
}
Self(out)
}

/// Lane-wise equality as a packed 16-bit bitmask.
///
/// Bit `i` of the result is set iff `self.lane(i) == other.lane(i)`. Bit
/// order is **LSB-first**: lane `0` occupies bit `0`. Same convention as
/// [`I32x16::cmpge_zero_mask`] and [`I32x16::gt_bitmask`].
///
/// Edge cases: equality is exact bitwise comparison over the full 32-bit
/// range, so `u32::MAX` and `0` behave like any other value — no
/// saturation, wrapping, or signedness question arises.
///
/// Plain index loop over the array polyfill; see [`I32x16::gt_bitmask`]
/// for why no intrinsic override is earned.
#[inline(always)]
pub fn eq_bitmask(self, other: Self) -> u16 {
let mut mask = 0u16;
for i in 0..16 {
if self.0[i] == other.0[i] {
mask |= 1 << i;
}
}
mask
}
}

// 256-bit int lanes — scalar polyfills filling the gap surfaced by the
Expand Down Expand Up @@ -2307,6 +2330,34 @@ impl I32x16 {
}
mask
}

/// Lane-wise **signed** greater-than as a packed 16-bit bitmask.
///
/// Bit `i` of the result is set iff `self.lane(i) > other.lane(i)` under
/// two's-complement signed ordering. Bit order is **LSB-first**: lane `0`
/// occupies bit `0`. Same convention as [`Self::cmpge_zero_mask`].
///
/// Edge cases (all exact; no saturation, wrapping, or clamping):
/// * `i32::MIN` as the threshold is set for every lane strictly greater
/// than it, and clear for lanes equal to `i32::MIN`.
/// * `i32::MAX` as the threshold yields `0` — no `i32` exceeds it.
/// * Comparison is signed, *not* bit-pattern: `-1 > 0` is `false`.
///
/// Plain index loop over the array polyfill — the codegen oracle
/// (`.claude/knowledge/simd-codegen-oracle/`) measured that LLVM lowers
/// compare-and-pack-to-bitmask shapes of exactly this form to packed
/// compares plus a `vmovmsk`-class extraction, so no `unsafe` and no
/// `core::arch` intrinsic override is earned here.
#[inline(always)]
pub fn gt_bitmask(self, other: Self) -> u16 {
let mut mask = 0u16;
for i in 0..16 {
if self.0[i] > other.0[i] {
mask |= 1 << i;
}
}
mask
}
}
impl Mul for I32x16 {
type Output = Self;
Expand Down
47 changes: 47 additions & 0 deletions src/simd_avx512.rs
Original file line number Diff line number Diff line change
Expand Up @@ -943,6 +943,32 @@ impl I32x16 {
unsafe { _mm512_cmpge_epi32_mask(self.0, _mm512_setzero_si512()) }
}

/// Lane-wise **signed** greater-than as a packed 16-bit bitmask.
///
/// Bit `i` of the result is set iff `self.lane(i) > other.lane(i)` under
/// two's-complement signed ordering. Bit order is **LSB-first**: lane `0`
/// occupies bit `0`. Sibling of [`Self::cmpge_zero_mask`], which uses the
/// same convention.
///
/// Edge cases (all exact, no saturation or clamping anywhere):
/// * `i32::MIN > i32::MIN` → `false`; nothing is greater than `i32::MIN`
/// except strictly larger values, so `x.gt_bitmask(splat(i32::MIN))` is
/// set for every lane except those equal to `i32::MIN`.
/// * `i32::MAX` as the threshold yields `0` — no `i32` exceeds it.
/// * Negative operands compare as signed, *not* as bit patterns:
/// `-1 > 0` is `false` even though `0xFFFF_FFFF > 0` unsigned.
///
/// AVX-512 lowers this to a single `VPCMPGTD` into a `__mmask16`, which
/// *is* a `u16` — the packed bitmask is the hardware's native result, so
/// there is no extraction step to elide.
#[inline(always)]
pub fn gt_bitmask(self, other: Self) -> u16 {
// SAFETY: `Self` wraps a native `__m512i` and this impl block is
// compiled only under the `avx512f` dispatch arm, the same guarantee
// every other method on this type relies on.
unsafe { _mm512_cmpgt_epi32_mask(self.0, other.0) }
}

#[inline(always)]
pub fn simd_min(self, other: Self) -> Self {
Self(unsafe { _mm512_min_epi32(self.0, other.0) })
Expand Down Expand Up @@ -1592,6 +1618,27 @@ impl U32x16 {
unsafe { _mm512_reduce_add_epi32(self.0) as u32 }
}

/// Lane-wise equality as a packed 16-bit bitmask.
///
/// Bit `i` of the result is set iff `self.lane(i) == other.lane(i)`. Bit
/// order is **LSB-first**: lane `0` occupies bit `0`. Same convention as
/// [`I32x16::cmpge_zero_mask`] and [`I32x16::gt_bitmask`].
///
/// Edge cases: equality is exact bitwise comparison over the full 32-bit
/// range, so `u32::MAX` and `0` behave like any other value and there is
/// no saturation, wrapping, or signedness question to resolve — an `i32`
/// lane pattern compares identically if reinterpreted.
///
/// AVX-512 lowers this to a single `VPCMPEQD` into a `__mmask16`, which
/// *is* a `u16` — the packed bitmask is the hardware's native result.
#[inline(always)]
pub fn eq_bitmask(self, other: Self) -> u16 {
// SAFETY: `Self` wraps a native `__m512i` and this impl block is
// compiled only under the `avx512f` dispatch arm, the same guarantee
// every other method on this type relies on.
unsafe { _mm512_cmpeq_epu32_mask(self.0, other.0) }
}

/// Lane-wise left-rotate by `n` bits — the ARX rotate (matches
/// `u32::rotate_left`), the third ChaCha20/BLAKE-family primitive alongside
/// `Add` + `BitXor`. Single `VPROLVD` (AVX-512F variable rotate). The rotate
Expand Down
Loading
Loading