From 116316280f04b6b8318282989d3072261e11058c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 24 Aug 2026 21:34:13 +0000 Subject: [PATCH] Add SIMD primitive comparison kernels Moves the explicit AVX2 and AVX-512 primitive comparison kernels from #9547 into the handwritten comparison path, leaving the RowFn work in #9547 and #9548 untouched. Non-x86 targets and x86-64 CPUs without AVX2 keep the portable lane-kernel fallback. Full-word kernels cover every primitive type, comparison operator, and array/constant orientation, with scalar tail handling, split by lane width so each width reads on its own. Boundary and nullable coverage comes with them. Signed-off-by: Connor Tsui --- .../src/scalar_fn/fns/binary/compare/mod.rs | 14 +- .../{primitive.rs => primitive/mod.rs} | 32 +- .../fns/binary/compare/primitive/simd/mod.rs | 419 ++++++++++++++++++ .../binary/compare/primitive/simd/portable.rs | 7 + .../binary/compare/primitive/simd/tests.rs | 332 ++++++++++++++ .../compare/primitive/simd/x86/float_order.rs | 82 ++++ .../compare/primitive/simd/x86/lanes_16.rs | 162 +++++++ .../compare/primitive/simd/x86/lanes_32.rs | 159 +++++++ .../compare/primitive/simd/x86/lanes_64.rs | 169 +++++++ .../compare/primitive/simd/x86/lanes_8.rs | 130 ++++++ .../binary/compare/primitive/simd/x86/mod.rs | 22 + .../compare/primitive/simd/x86/shared.rs | 171 +++++++ .../fns/binary/compare/primitive/tests.rs | 197 ++++++++ 13 files changed, 1888 insertions(+), 8 deletions(-) rename vortex-array/src/scalar_fn/fns/binary/compare/{primitive.rs => primitive/mod.rs} (84%) create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/mod.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/portable.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/tests.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/float_order.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_16.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_32.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_64.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_8.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/mod.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/shared.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/tests.rs diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index 8b5a4522b07..62474dfdf94 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -4,9 +4,9 @@ //! Native comparison kernels. //! //! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every -//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from -//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise -//! comparator for nested types. There is no Arrow fallback. +//! comparison directly over Vortex canonical arrays: bit buffers for booleans, explicit x86 SIMD +//! or portable lane kernels for primitives, lane kernels for decimals, binary views for +//! strings/bytes, and a row-wise comparator for nested types. There is no Arrow fallback. //! //! Floating point values compare with Vortex's total ordering (`NaN` is the largest value, //! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics. @@ -282,8 +282,14 @@ pub(super) fn ordering_predicate(op: CompareOperator) -> fn(Ordering) -> bool { } /// Freeze `len` bits packed into `words` (LSB-first, 64 lanes per word) into a [`BitBuffer`]. -pub(super) fn bit_buffer_from_words(words: BufferMut, len: usize) -> BitBuffer { +pub(super) fn bit_buffer_from_words(mut words: BufferMut, len: usize) -> BitBuffer { debug_assert!(words.len() * 64 >= len); + + // Byte reinterpretation must preserve LSB-first lane order on big-endian hosts. + for word in words.iter_mut() { + *word = word.to_le(); + } + let mut bytes = words.into_byte_buffer(); bytes.truncate(len.div_ceil(8)); BitBuffer::new(bytes.freeze(), len) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/mod.rs similarity index 84% rename from vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs rename to vortex-array/src/scalar_fn/fns/binary/compare/primitive/mod.rs index 1247358dce1..5584835bbef 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/mod.rs @@ -1,7 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays via bit-packing lane kernels. +//! Native comparison of primitive arrays with Vortex's scalar ordering semantics. +//! +//! The outer comparison path owns decoding, validity, and the portable bit-packing fallback. +//! On x86-64, [`simd`] can replace the value loop with explicit AVX2 or AVX-512 kernels. + +mod simd; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -41,7 +46,7 @@ pub(super) fn compare_primitive( }) } -fn compare_primitive_typed( +fn compare_primitive_typed( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, @@ -105,7 +110,15 @@ fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { } } -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { +fn compare_slices( + lhs: &[T], + rhs: &[T], + op: CompareOperator, +) -> BitBuffer { + if let Some(bits) = simd::try_compare_slices(lhs, rhs, op) { + return bits; + } + // Dispatch the operator outside the lane loop so each instantiation vectorizes a single // branch-free predicate. match op { @@ -118,7 +131,15 @@ fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> } } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { +fn compare_slice_constant( + lhs: &[T], + rhs: T, + op: CompareOperator, +) -> BitBuffer { + if let Some(bits) = simd::try_compare_slice_constant(lhs, rhs, op) { + return bits; + } + match op { CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), @@ -128,3 +149,6 @@ fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), } } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/mod.rs new file mode 100644 index 00000000000..6b9ca6aa5f3 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/mod.rs @@ -0,0 +1,419 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Explicit SIMD primitive comparison kernels that write bitmap words directly. +//! +//! Safe entry points validate input and output bounds before invoking AVX2 or AVX-512. The +//! explicit kernels handle complete bitmap words, while a scalar tail handles remaining rows. +//! Other architectures and x86-64 CPUs without AVX2 return to the caller's portable fallback. + +#[cfg(not(target_arch = "x86_64"))] +mod portable; +#[cfg(target_arch = "x86_64")] +mod x86; + +use vortex_buffer::BitBuffer; +#[cfg(target_arch = "x86_64")] +use vortex_buffer::BufferMut; + +use crate::dtype::NativePType; +#[cfg(target_arch = "x86_64")] +use crate::scalar_fn::fns::binary::compare::bit_buffer_from_words; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare two slices with the widest available explicit SIMD kernel. +/// +/// Returns `None` outside x86-64 and when the current x86-64 CPU does not support AVX2. +#[cfg(target_arch = "x86_64")] +pub(super) fn try_compare_slices( + lhs: &[T], + rhs: &[T], + op: CompareOperator, +) -> Option { + assert!( + lhs.len() == rhs.len(), + "primitive comparison inputs must have equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + + try_compare_inputs( + PrimitiveInput::Slice(lhs), + PrimitiveInput::Slice(rhs), + op, + lhs.len(), + ) +} + +#[cfg(not(target_arch = "x86_64"))] +/// Return control to the portable primitive comparison path outside x86-64. +pub(super) fn try_compare_slices( + _lhs: &[T], + _rhs: &[T], + _op: CompareOperator, +) -> Option { + None +} + +/// Compare a slice with a constant using the widest available explicit SIMD kernel. +/// +/// Returns `None` outside x86-64 and when the current x86-64 CPU does not support AVX2. +#[cfg(target_arch = "x86_64")] +pub(super) fn try_compare_slice_constant( + values: &[T], + constant: T, + op: CompareOperator, +) -> Option { + try_compare_inputs( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + op, + values.len(), + ) +} + +#[cfg(not(target_arch = "x86_64"))] +/// Return control to the portable primitive comparison path outside x86-64. +pub(super) fn try_compare_slice_constant( + _values: &[T], + _constant: T, + _op: CompareOperator, +) -> Option { + None +} + +#[cfg(target_arch = "x86_64")] +fn try_compare_inputs( + lhs: PrimitiveInput<'_, T>, + rhs: PrimitiveInput<'_, T>, + op: CompareOperator, + row_count: usize, +) -> Option { + let detected = DetectedSimd::detect::()?; + let mut words = BufferMut::::zeroed(row_count.div_ceil(64)); + + compare_into_words(detected, lhs, rhs, op, row_count, words.as_mut_slice()); + + Some(bit_buffer_from_words(words, row_count)) +} + +#[cfg(target_arch = "x86_64")] +#[derive(Clone, Copy)] +enum PrimitiveInput<'a, T> { + Slice(&'a [T]), + Constant(T), +} + +#[cfg(target_arch = "x86_64")] +impl<'a, T: NativePType> PrimitiveInput<'a, T> { + fn assert_addresses_rows(self, side: &str, row_count: usize) { + if let Self::Slice(values) = self { + assert!( + values.len() >= row_count, + "{side} primitive comparison input must cover {row_count} rows, got {}", + values.len() + ); + } + } + + /// Read one logical input row without checking its index. + /// + /// # Safety + /// + /// For [`Self::Slice`], `index` must be less than the slice length. A constant accepts every + /// logical row index. Violating the slice requirement causes undefined behavior. + unsafe fn get_unchecked(self, index: usize) -> T { + match self { + // SAFETY: forwarded from this method's contract. + Self::Slice(values) => unsafe { *values.get_unchecked(index) }, + Self::Constant(value) => value, + } + } +} + +#[cfg(target_arch = "x86_64")] +fn compare_into_words( + detected: DetectedSimd, + lhs: PrimitiveInput<'_, T>, + rhs: PrimitiveInput<'_, T>, + op: CompareOperator, + row_count: usize, + words: &mut [u64], +) { + let required_words = row_count.div_ceil(64); + assert!( + words.len() >= required_words, + "primitive comparison output must contain {required_words} words, got {}", + words.len() + ); + lhs.assert_addresses_rows("left", row_count); + rhs.assert_addresses_rows("right", row_count); + + let full_words = row_count / 64; + let comparison = DenseComparison::new(lhs, rhs, op); + detected.compare(comparison, &mut words[..full_words]); + + let remainder = row_count % 64; + if remainder != 0 { + let base = full_words * 64; + words[full_words] = (0..remainder).fold(0, |packed, bit| { + // SAFETY: `base + bit < row_count` for every tail lane. + let lhs = unsafe { lhs.get_unchecked(base + bit) }; + // SAFETY: `base + bit < row_count` for every tail lane. + let rhs = unsafe { rhs.get_unchecked(base + bit) }; + packed | ((super::apply_op(lhs, rhs, op) as u64) << bit) + }); + } +} + +#[cfg(target_arch = "x86_64")] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum SimdLevel { + /// AVX-512F, plus AVX-512BW when the physical type requires it. + Avx512, + + /// AVX2. + Avx2, +} + +/// Proof that the current CPU supports one explicit SIMD implementation for a physical type. +/// +/// Only the detection methods construct this token. An AVX-512 token also proves AVX-512BW support +/// when [`SimdCompare::REQUIRES_AVX512BW`] is true. +#[cfg(target_arch = "x86_64")] +#[derive(Clone, Copy)] +struct DetectedSimd(SimdLevel); + +#[cfg(target_arch = "x86_64")] +impl DetectedSimd { + fn detect() -> Option { + // Benchmark builds select from compile-time target features so the selected ISA is visible + // to LLVM. Production builds retain runtime detection before any target-feature call. + #[cfg(feature = "_test-harness")] + let avx512_available = cfg!(target_feature = "avx512f") + && (!T::REQUIRES_AVX512BW || cfg!(target_feature = "avx512bw")); + #[cfg(not(feature = "_test-harness"))] + let avx512_available = std::arch::is_x86_feature_detected!("avx512f") + && (!T::REQUIRES_AVX512BW || std::arch::is_x86_feature_detected!("avx512bw")); + + if avx512_available { + return Some(Self(SimdLevel::Avx512)); + } + + #[cfg(feature = "_test-harness")] + let avx2_available = cfg!(target_feature = "avx2"); + #[cfg(not(feature = "_test-harness"))] + let avx2_available = std::arch::is_x86_feature_detected!("avx2"); + + avx2_available.then_some(Self(SimdLevel::Avx2)) + } + + #[cfg(test)] + fn detect_runtime() -> Option { + let avx512_available = std::arch::is_x86_feature_detected!("avx512f") + && (!T::REQUIRES_AVX512BW || std::arch::is_x86_feature_detected!("avx512bw")); + if avx512_available { + return Some(Self(SimdLevel::Avx512)); + } + + std::arch::is_x86_feature_detected!("avx2").then_some(Self(SimdLevel::Avx2)) + } + + fn compare(self, comparison: DenseComparison<'_, T>, words: &mut [u64]) { + match self.0 { + SimdLevel::Avx512 => { + // SAFETY: the detection methods prove the type's required AVX-512 features. + // `compare_into_words` proves every slice covers these words. + unsafe { T::compare_words_avx512(comparison, words) } + } + SimdLevel::Avx2 => { + // SAFETY: the detection methods prove AVX2 support. + // `compare_into_words` proves every slice covers these words. + unsafe { T::compare_words_avx2(comparison, words) } + } + } + } +} + +/// An opaque, normalized comparison shape for explicit SIMD kernels. +#[cfg(target_arch = "x86_64")] +#[derive(Clone, Copy)] +pub(super) struct DenseComparison<'a, T> { + kind: DenseComparisonKind<'a, T>, +} + +#[cfg(target_arch = "x86_64")] +#[derive(Clone, Copy)] +enum DenseComparisonKind<'a, T> { + EqualArrays { + lhs: &'a [T], + rhs: &'a [T], + invert: bool, + }, + GreaterArrays { + lhs: &'a [T], + rhs: &'a [T], + invert: bool, + }, + EqualArrayConstant { + values: &'a [T], + constant: T, + invert: bool, + }, + GreaterArrayConstant { + values: &'a [T], + constant: T, + invert: bool, + }, + GreaterConstantArray { + constant: T, + values: &'a [T], + invert: bool, + }, +} + +#[cfg(target_arch = "x86_64")] +impl<'a, T: NativePType> DenseComparison<'a, T> { + fn new(lhs: PrimitiveInput<'a, T>, rhs: PrimitiveInput<'a, T>, op: CompareOperator) -> Self { + if let (PrimitiveInput::Constant(lhs), PrimitiveInput::Slice(rhs)) = (lhs, rhs) { + return Self::new( + PrimitiveInput::Slice(rhs), + PrimitiveInput::Constant(lhs), + op.swap(), + ); + } + + let kind = match (lhs, rhs, op) { + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Eq) => { + DenseComparisonKind::EqualArrays { + lhs, + rhs, + invert: false, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::NotEq) => { + DenseComparisonKind::EqualArrays { + lhs, + rhs, + invert: true, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Gt) => { + DenseComparisonKind::GreaterArrays { + lhs, + rhs, + invert: false, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Gte) => { + DenseComparisonKind::GreaterArrays { + lhs: rhs, + rhs: lhs, + invert: true, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Lt) => { + DenseComparisonKind::GreaterArrays { + lhs: rhs, + rhs: lhs, + invert: false, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Lte) => { + DenseComparisonKind::GreaterArrays { + lhs, + rhs, + invert: true, + } + } + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Eq, + ) => DenseComparisonKind::EqualArrayConstant { + values, + constant, + invert: false, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::NotEq, + ) => DenseComparisonKind::EqualArrayConstant { + values, + constant, + invert: true, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Gt, + ) => DenseComparisonKind::GreaterArrayConstant { + values, + constant, + invert: false, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Gte, + ) => DenseComparisonKind::GreaterConstantArray { + constant, + values, + invert: true, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Lt, + ) => DenseComparisonKind::GreaterConstantArray { + constant, + values, + invert: false, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Lte, + ) => DenseComparisonKind::GreaterArrayConstant { + values, + constant, + invert: true, + }, + (PrimitiveInput::Constant(_), PrimitiveInput::Constant(_), _) => unreachable!(), + (PrimitiveInput::Constant(_), PrimitiveInput::Slice(_), _) => unreachable!(), + }; + + Self { kind } + } +} + +/// Primitive types supported by explicit x86-64 kernels or the portable fallback. +pub(super) trait SimdCompare: NativePType { + /// Whether this type's AVX-512 implementation also requires AVX-512BW. + #[cfg(target_arch = "x86_64")] + const REQUIRES_AVX512BW: bool; + + /// Compare complete 64-row chunks with AVX2 and write one bitmap word per chunk. + /// + /// # Safety + /// + /// The current CPU must support AVX2. Every slice in `comparison` must contain at least + /// `words.len() * 64` elements. Violating either requirement causes undefined behavior. + #[cfg(target_arch = "x86_64")] + unsafe fn compare_words_avx2(comparison: DenseComparison<'_, Self>, words: &mut [u64]); + + /// Compare complete 64-row chunks with AVX-512 and write one bitmap word per chunk. + /// + /// # Safety + /// + /// The current CPU must support AVX-512F and, when [`Self::REQUIRES_AVX512BW`] is true, + /// AVX-512BW. Every slice in `comparison` must contain at least `words.len() * 64` elements. + /// Violating any requirement causes undefined behavior. + #[cfg(target_arch = "x86_64")] + unsafe fn compare_words_avx512(comparison: DenseComparison<'_, Self>, words: &mut [u64]); +} + +#[cfg(all(test, target_arch = "x86_64"))] +#[allow(clippy::cast_possible_truncation)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/portable.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/portable.rs new file mode 100644 index 00000000000..7b3ccd525b4 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/portable.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use super::SimdCompare; +use crate::dtype::NativePType; + +impl SimdCompare for T {} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/tests.rs new file mode 100644 index 00000000000..bfade69463c --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/tests.rs @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Debug; + +use half::f16; + +use super::DenseComparison; +use super::DetectedSimd; +use super::PrimitiveInput; +use super::SimdCompare; +use super::SimdLevel; +use super::compare_into_words; +use crate::dtype::NativePType; +use crate::scalar_fn::fns::operators::CompareOperator; + +const OPS: [CompareOperator; 6] = [ + CompareOperator::Eq, + CompareOperator::NotEq, + CompareOperator::Gt, + CompareOperator::Gte, + CompareOperator::Lt, + CompareOperator::Lte, +]; + +#[cfg(feature = "_test-harness")] +#[test] +fn test_benchmark_dispatch_uses_compiled_features() { + let expected_i64 = if cfg!(target_feature = "avx512f") { + Some(SimdLevel::Avx512) + } else if cfg!(target_feature = "avx2") { + Some(SimdLevel::Avx2) + } else { + None + }; + let expected_i8 = if cfg!(target_feature = "avx512f") && cfg!(target_feature = "avx512bw") { + Some(SimdLevel::Avx512) + } else if cfg!(target_feature = "avx2") { + Some(SimdLevel::Avx2) + } else { + None + }; + + assert_eq!( + DetectedSimd::detect::().map(|detected| detected.0), + expected_i64 + ); + assert_eq!( + DetectedSimd::detect::().map(|detected| detected.0), + expected_i8 + ); +} + +fn input_at(input: PrimitiveInput<'_, T>, index: usize) -> T { + match input { + PrimitiveInput::Slice(values) => values[index], + PrimitiveInput::Constant(value) => value, + } +} + +fn apply_reference(lhs: T, rhs: T, op: CompareOperator) -> bool { + match op { + CompareOperator::Eq => lhs.is_eq(rhs), + CompareOperator::NotEq => !lhs.is_eq(rhs), + CompareOperator::Gt => lhs.is_gt(rhs), + CompareOperator::Gte => lhs.is_ge(rhs), + CompareOperator::Lt => lhs.is_lt(rhs), + CompareOperator::Lte => lhs.is_le(rhs), + } +} + +fn expected_word( + lhs: PrimitiveInput<'_, T>, + rhs: PrimitiveInput<'_, T>, + op: CompareOperator, +) -> u64 { + (0..64).fold(0, |word, bit| { + let lhs = input_at(lhs, bit); + let rhs = input_at(rhs, bit); + + word | ((apply_reference(lhs, rhs, op) as u64) << bit) + }) +} + +fn assert_simd_matches_scalar(lhs: &[T; 64], rhs: &[T; 64]) +where + T: SimdCompare + Debug, +{ + let shapes = [ + ( + "array-array", + PrimitiveInput::Slice(lhs), + PrimitiveInput::Slice(rhs), + ), + ( + "array-constant", + PrimitiveInput::Slice(lhs), + PrimitiveInput::Constant(rhs[7]), + ), + ( + "constant-array", + PrimitiveInput::Constant(lhs[11]), + PrimitiveInput::Slice(rhs), + ), + ]; + + for op in OPS { + for (shape, lhs, rhs) in shapes { + let expected = expected_word(lhs, rhs, op); + let comparison = DenseComparison::new(lhs, rhs, op); + if std::arch::is_x86_feature_detected!("avx2") { + // SAFETY: runtime detection proves AVX2 support, and the inputs cover 64 rows. + let mut actual = [0]; + unsafe { T::compare_words_avx2(comparison, &mut actual) }; + assert_eq!(actual[0], expected, "AVX2 shape={shape} op={op:?}"); + } + if std::arch::is_x86_feature_detected!("avx512f") + && (!T::REQUIRES_AVX512BW || std::arch::is_x86_feature_detected!("avx512bw")) + { + // SAFETY: runtime detection proves the type's required AVX-512 features, and + // the inputs cover 64 rows. + let mut actual = [0]; + unsafe { T::compare_words_avx512(comparison, &mut actual) }; + assert_eq!(actual[0], expected, "AVX-512 shape={shape} op={op:?}"); + } + } + } + + assert_word_loop_matches_scalar(lhs, rhs); +} + +fn assert_word_loop_matches_scalar(lhs_seed: &[T; 64], rhs_seed: &[T; 64]) +where + T: SimdCompare + Debug, +{ + let Some(detected) = DetectedSimd::detect_runtime::() else { + return; + }; + + for len in [ + 0, // Empty input. + 1, // One scalar tail lane. + 7, // One lane before an eight-row mask chunk. + 8, // One complete eight-row mask chunk. + 9, // One lane after an eight-row mask chunk. + 63, // One lane before a complete bitmap word. + 64, // One complete bitmap word. + 65, // One lane after a complete bitmap word. + 129, // Two complete bitmap words and one tail lane. + ] { + let lhs_values = lhs_seed + .iter() + .copied() + .cycle() + .take(len) + .collect::>(); + let rhs_values = rhs_seed + .iter() + .copied() + .cycle() + .take(len) + .collect::>(); + let shapes = [ + ( + "array-array", + PrimitiveInput::Slice(lhs_values.as_slice()), + PrimitiveInput::Slice(rhs_values.as_slice()), + ), + ( + "array-constant", + PrimitiveInput::Slice(lhs_values.as_slice()), + PrimitiveInput::Constant(rhs_seed[7]), + ), + ( + "constant-array", + PrimitiveInput::Constant(lhs_seed[11]), + PrimitiveInput::Slice(rhs_values.as_slice()), + ), + ]; + + for op in OPS { + for (shape, lhs, rhs) in shapes { + let mut actual = vec![u64::MAX; len.div_ceil(64)]; + compare_into_words(detected, lhs, rhs, op, len, &mut actual); + + let mut expected = vec![0; len.div_ceil(64)]; + for index in 0..len { + let lhs = input_at(lhs, index); + let rhs = input_at(rhs, index); + expected[index / 64] |= (apply_reference(lhs, rhs, op) as u64) << (index % 64); + } + + assert_eq!( + actual, expected, + "word loop ISA={:?} shape={shape} len={len} op={op:?}", + detected.0 + ); + } + } + } +} + +#[test] +fn test_u8_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as u8).wrapping_mul(37)); + let rhs = std::array::from_fn(|index| (index as u8).wrapping_mul(19).wrapping_add(127)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_i8_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as i8).wrapping_mul(37)); + let rhs = std::array::from_fn(|index| (index as i8).wrapping_mul(-19).wrapping_add(63)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_signed_16_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as i16).wrapping_mul(i16::MAX / 31)); + let rhs = std::array::from_fn(|index| (index as i16).wrapping_mul(i16::MIN / 29)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_unsigned_16_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as u16).wrapping_mul(u16::MAX / 31)); + let rhs = std::array::from_fn(|index| (index as u16).wrapping_mul(u16::MAX / 29)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_signed_32_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as i32).wrapping_mul(i32::MAX / 31)); + let rhs = std::array::from_fn(|index| (index as i32).wrapping_mul(i32::MIN / 29)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_unsigned_32_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as u32).wrapping_mul(u32::MAX / 31)); + let rhs = std::array::from_fn(|index| (index as u32).wrapping_mul(u32::MAX / 29)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_signed_64_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as i64).wrapping_mul(i64::MAX / 31)); + let rhs = std::array::from_fn(|index| (index as i64).wrapping_mul(i64::MIN / 29)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_unsigned_64_masks_match_scalar() { + let lhs = std::array::from_fn(|index| (index as u64).wrapping_mul(u64::MAX / 31)); + let rhs = std::array::from_fn(|index| (index as u64).wrapping_mul(u64::MAX / 29)); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_float_64_total_order_masks_match_scalar() { + const BITS: [u64; 16] = [ + 0xfff8_0000_0000_0001, // Negative quiet NaN with a payload. + 0xfff0_0000_0000_0000, // Negative infinity. + 0xbff0_0000_0000_0000, // Negative one. + 0x8000_0000_0000_0001, // Negative minimum subnormal. + 0x8000_0000_0000_0000, // Negative zero. + 0x0000_0000_0000_0000, // Positive zero. + 0x0000_0000_0000_0001, // Positive minimum subnormal. + 0x3ff0_0000_0000_0000, // Positive one. + 0x7ff0_0000_0000_0000, // Positive infinity. + 0x7ff8_0000_0000_0000, // Positive quiet NaN. + 0x7ff8_0000_0000_0001, // Positive quiet NaN with a payload. + 0x7fff_ffff_ffff_ffff, // Positive NaN with the maximum payload. + 0xfff8_0000_0000_0000, // Negative quiet NaN. + 0x4000_0000_0000_0000, // Positive two. + 0xc000_0000_0000_0000, // Negative two. + 0x3fe0_0000_0000_0000, // Positive one half. + ]; + let lhs = std::array::from_fn(|index| f64::from_bits(BITS[index % BITS.len()])); + let rhs = std::array::from_fn(|index| f64::from_bits(BITS[(index * 7 + 3) % BITS.len()])); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_float_16_total_order_masks_match_scalar() { + const BITS: [u16; 16] = [ + 0xfe01, // Negative quiet NaN with a payload. + 0xfc00, // Negative infinity. + 0xbc00, // Negative one. + 0x8001, // Negative minimum subnormal. + 0x8000, // Negative zero. + 0x0000, // Positive zero. + 0x0001, // Positive minimum subnormal. + 0x3c00, // Positive one. + 0x7c00, // Positive infinity. + 0x7e00, // Positive quiet NaN. + 0x7e01, // Positive quiet NaN with a payload. + 0x7fff, // Positive NaN with the maximum payload. + 0xfe00, // Negative quiet NaN. + 0x4000, // Positive two. + 0xc000, // Negative two. + 0x3800, // Positive one half. + ]; + let lhs = std::array::from_fn(|index| f16::from_bits(BITS[index % BITS.len()])); + let rhs = std::array::from_fn(|index| f16::from_bits(BITS[(index * 7 + 3) % BITS.len()])); + assert_simd_matches_scalar(&lhs, &rhs); +} + +#[test] +fn test_float_32_total_order_masks_match_scalar() { + const BITS: [u32; 16] = [ + 0xffc0_0001, // Negative quiet NaN with a payload. + 0xff80_0000, // Negative infinity. + 0xbf80_0000, // Negative one. + 0x8000_0001, // Negative minimum subnormal. + 0x8000_0000, // Negative zero. + 0x0000_0000, // Positive zero. + 0x0000_0001, // Positive minimum subnormal. + 0x3f80_0000, // Positive one. + 0x7f80_0000, // Positive infinity. + 0x7fc0_0000, // Positive quiet NaN. + 0x7fc0_0001, // Positive quiet NaN with a payload. + 0x7fff_ffff, // Positive NaN with the maximum payload. + 0xffc0_0000, // Negative quiet NaN. + 0x4000_0000, // Positive two. + 0xc000_0000, // Negative two. + 0x3f00_0000, // Positive one half. + ]; + let lhs = std::array::from_fn(|index| f32::from_bits(BITS[index % BITS.len()])); + let rhs = std::array::from_fn(|index| f32::from_bits(BITS[(index * 7 + 3) % BITS.len()])); + assert_simd_matches_scalar(&lhs, &rhs); +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/float_order.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/float_order.rs new file mode 100644 index 00000000000..bf9c88cf52f --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/float_order.rs @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::*; + +/// Convert f16 bits into signed integer keys with the same total ordering. +/// +/// # Safety +/// +/// The current CPU must support AVX2. Calling this function without AVX2 support causes +/// undefined behavior. +#[target_feature(enable = "avx2")] +#[inline] +pub(super) unsafe fn float_key_16(bits: __m256i) -> __m256i { + let negative = _mm256_srai_epi16::<15>(bits); + _mm256_xor_si256(bits, _mm256_srli_epi16::<1>(negative)) +} + +/// Convert f16 bits into signed integer keys with the same total ordering. +/// +/// # Safety +/// +/// The current CPU must support AVX-512F and AVX-512BW. Calling this function without both +/// features causes undefined behavior. +#[target_feature(enable = "avx512f,avx512bw")] +#[inline] +pub(super) unsafe fn float_key_16_avx512(bits: __m512i) -> __m512i { + let negative = _mm512_srai_epi16::<15>(bits); + _mm512_xor_si512(bits, _mm512_srli_epi16::<1>(negative)) +} + +/// Convert f32 bits into signed integer keys with the same total ordering. +/// +/// # Safety +/// +/// The current CPU must support AVX2. Calling this function without AVX2 support causes +/// undefined behavior. +#[target_feature(enable = "avx2")] +#[inline] +pub(super) unsafe fn float_key_32(bits: __m256i) -> __m256i { + let negative = _mm256_srai_epi32::<31>(bits); + _mm256_xor_si256(bits, _mm256_srli_epi32::<1>(negative)) +} + +/// Convert f32 bits into signed integer keys with the same total ordering. +/// +/// # Safety +/// +/// The current CPU must support AVX-512F. Calling this function without AVX-512F support +/// causes undefined behavior. +#[target_feature(enable = "avx512f")] +#[inline] +pub(super) unsafe fn float_key_32_avx512(bits: __m512i) -> __m512i { + let negative = _mm512_srai_epi32::<31>(bits); + _mm512_xor_si512(bits, _mm512_srli_epi32::<1>(negative)) +} + +/// Convert f64 bits into signed integer keys with the same total ordering. +/// +/// # Safety +/// +/// The current CPU must support AVX2. Calling this function without AVX2 support causes +/// undefined behavior. +#[target_feature(enable = "avx2")] +#[inline] +pub(super) unsafe fn float_key_64(bits: __m256i) -> __m256i { + let negative = _mm256_cmpgt_epi64(_mm256_setzero_si256(), bits); + _mm256_xor_si256(bits, _mm256_srli_epi64::<1>(negative)) +} + +/// Convert f64 bits into signed integer keys with the same total ordering. +/// +/// # Safety +/// +/// The current CPU must support AVX-512F. Calling this function without AVX-512F support +/// causes undefined behavior. +#[target_feature(enable = "avx512f")] +#[inline] +pub(super) unsafe fn float_key_64_avx512(bits: __m512i) -> __m512i { + let negative = _mm512_srai_epi64::<63>(bits); + _mm512_xor_si512(bits, _mm512_srli_epi64::<1>(negative)) +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_16.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_16.rs new file mode 100644 index 00000000000..4a6afb137bf --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_16.rs @@ -0,0 +1,162 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::*; + +use half::f16; + +use super::super::DenseComparison; +use super::super::DenseComparisonKind; +use super::super::SimdCompare; +use super::float_order::float_key_16; +use super::float_order::float_key_16_avx512; +use super::shared::ToSignedBits; +use super::shared::compact_even_bits; + +const GREATER_THAN: i32 = _MM_CMPINT_NLE; + +macro_rules! impl_16 { + ($t:ty, $unsigned:expr, $float:expr) => { + #[allow(clippy::cast_possible_truncation)] + impl SimdCompare for $t { + const REQUIRES_AVX512BW: bool = true; + + #[target_feature(enable = "avx2")] + unsafe fn compare_words_avx2( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 4) { + let offset = chunk * 16; + let lhs = load_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi16, + __m256i + ); + let rhs = load_or_broadcast!( + $rhs, + offset, + $rhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi16, + __m256i + ); + let mask = if $eq { + compact_even_bits( + _mm256_movemask_epi8(_mm256_cmpeq_epi16(lhs, rhs)) as u32, + ) as u16 + } else { + let (lhs, rhs) = if $float { + // SAFETY: this method's target-feature contract includes + // AVX2. + unsafe { (float_key_16(lhs), float_key_16(rhs)) } + } else if $unsigned { + let sign = _mm256_set1_epi16(i16::MIN); + ( + _mm256_xor_si256(lhs, sign), + _mm256_xor_si256(rhs, sign), + ) + } else { + (lhs, rhs) + }; + compact_even_bits( + _mm256_movemask_epi8(_mm256_cmpgt_epi16(lhs, rhs)) as u32, + ) as u16 + }; + + // SAFETY: each output value represents 16 rows. The caller + // provides four values for each complete 64-row output word. + unsafe { output.add(chunk).write(mask ^ ($invert as u16)) }; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn compare_words_avx512( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 2) { + let offset = chunk * 32; + let lhs = load_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi16, + __m512i + ); + let rhs = load_or_broadcast!( + $rhs, + offset, + $rhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi16, + __m512i + ); + let mask = if $eq { + _mm512_cmpeq_epi16_mask(lhs, rhs) + } else { + let (lhs, rhs) = if $float { + // SAFETY: this method's target-feature contract includes + // AVX-512BW. + unsafe { + ( + float_key_16_avx512(lhs), + float_key_16_avx512(rhs), + ) + } + } else { + (lhs, rhs) + }; + if $unsigned { + _mm512_cmp_epu16_mask::(lhs, rhs) + } else { + _mm512_cmp_epi16_mask::(lhs, rhs) + } + }; + + // SAFETY: each output value represents 32 rows. The caller + // provides two values for each complete 64-row output word. + unsafe { output.add(chunk).write(mask ^ ($invert as u32)) }; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + } + }; +} + +impl_16!(i16, false, false); +impl_16!(u16, true, false); +impl_16!(f16, false, true); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_32.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_32.rs new file mode 100644 index 00000000000..9f36d6557d0 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_32.rs @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::*; + +use super::super::DenseComparison; +use super::super::DenseComparisonKind; +use super::super::SimdCompare; +use super::float_order::float_key_32; +use super::float_order::float_key_32_avx512; +use super::shared::ToSignedBits; + +const GREATER_THAN: i32 = _MM_CMPINT_NLE; + +macro_rules! impl_32 { + ($t:ty, $unsigned:expr, $float:expr) => { + #[allow(clippy::cast_possible_truncation)] + impl SimdCompare for $t { + const REQUIRES_AVX512BW: bool = false; + + #[target_feature(enable = "avx2")] + unsafe fn compare_words_avx2( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 8) { + let offset = chunk * 8; + let lhs = load_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi32, + __m256i + ); + let rhs = load_or_broadcast!( + $rhs, + offset, + $rhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi32, + __m256i + ); + let mask = if $eq { + _mm256_movemask_ps(_mm256_castsi256_ps(_mm256_cmpeq_epi32( + lhs, rhs, + ))) as u8 + } else { + let (lhs, rhs) = if $float { + // SAFETY: this method's target-feature contract includes + // AVX2. + unsafe { (float_key_32(lhs), float_key_32(rhs)) } + } else if $unsigned { + let sign = _mm256_set1_epi32(i32::MIN); + ( + _mm256_xor_si256(lhs, sign), + _mm256_xor_si256(rhs, sign), + ) + } else { + (lhs, rhs) + }; + _mm256_movemask_ps(_mm256_castsi256_ps(_mm256_cmpgt_epi32( + lhs, rhs, + ))) as u8 + }; + + // SAFETY: each output byte represents eight rows. The caller + // provides eight bytes for each complete 64-row output word. + unsafe { output.add(chunk).write(mask ^ ($invert as u8)) }; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + + #[target_feature(enable = "avx512f")] + unsafe fn compare_words_avx512( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 4) { + let offset = chunk * 16; + let lhs = load_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi32, + __m512i + ); + let rhs = load_or_broadcast!( + $rhs, + offset, + $rhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi32, + __m512i + ); + let mask = if $eq { + _mm512_cmpeq_epi32_mask(lhs, rhs) + } else { + let (lhs, rhs) = if $float { + // SAFETY: this method's target-feature contract includes + // AVX-512F. + unsafe { + ( + float_key_32_avx512(lhs), + float_key_32_avx512(rhs), + ) + } + } else { + (lhs, rhs) + }; + if $unsigned { + _mm512_cmp_epu32_mask::(lhs, rhs) + } else { + _mm512_cmp_epi32_mask::(lhs, rhs) + } + }; + + // SAFETY: each output value represents 16 rows. The caller + // provides four values for each complete 64-row output word. + unsafe { output.add(chunk).write(mask ^ ($invert as u16)) }; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + } + }; +} + +impl_32!(i32, false, false); +impl_32!(u32, true, false); +impl_32!(f32, false, true); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_64.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_64.rs new file mode 100644 index 00000000000..c7ef84550c6 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_64.rs @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::*; + +use super::super::DenseComparison; +use super::super::DenseComparisonKind; +use super::super::SimdCompare; +use super::float_order::float_key_64; +use super::float_order::float_key_64_avx512; +use super::shared::ToSignedBits; + +const GREATER_THAN: i32 = _MM_CMPINT_NLE; + +macro_rules! impl_64 { + ($t:ty, $unsigned:expr, $float:expr) => { + #[allow(clippy::cast_possible_truncation)] + impl SimdCompare for $t { + const REQUIRES_AVX512BW: bool = false; + + #[target_feature(enable = "avx2")] + unsafe fn compare_words_avx2( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + macro_rules! compare_chunk { + ($offset:expr) => {{ + let lhs = load_or_broadcast!( + $lhs, + $offset, + $lhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi64x, + __m256i + ); + let rhs = load_or_broadcast!( + $rhs, + $offset, + $rhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi64x, + __m256i + ); + let mask = if $eq { + (_mm256_movemask_pd(_mm256_castsi256_pd( + _mm256_cmpeq_epi64(lhs, rhs), + )) as u8) + & 0xf + } else { + let (lhs, rhs) = if $float { + // SAFETY: the target-feature contract includes AVX2. + unsafe { (float_key_64(lhs), float_key_64(rhs)) } + } else if $unsigned { + let sign = _mm256_set1_epi64x(i64::MIN); + ( + _mm256_xor_si256(lhs, sign), + _mm256_xor_si256(rhs, sign), + ) + } else { + (lhs, rhs) + }; + (_mm256_movemask_pd(_mm256_castsi256_pd( + _mm256_cmpgt_epi64(lhs, rhs), + )) as u8) + & 0xf + }; + + mask + }}; + } + + let output = $words.as_mut_ptr().cast::(); + for byte in 0..($words.len() * 8) { + let offset = byte * 8; + let low = compare_chunk!(offset); + let high = compare_chunk!(offset + 4); + let mask = (low | (high << 4)) ^ ($invert as u8); + + // SAFETY: each output byte represents eight rows. The caller + // provides eight bytes for each complete 64-row output word. + unsafe { output.add(byte).write(mask) }; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + + #[target_feature(enable = "avx512f")] + unsafe fn compare_words_avx512( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 8) { + let offset = chunk * 8; + let lhs = load_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi64, + __m512i + ); + let rhs = load_or_broadcast!( + $rhs, + offset, + $rhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi64, + __m512i + ); + let mask = if $eq { + _mm512_cmpeq_epi64_mask(lhs, rhs) + } else { + let (lhs, rhs) = if $float { + // SAFETY: the target-feature contract includes AVX-512F. + unsafe { + ( + float_key_64_avx512(lhs), + float_key_64_avx512(rhs), + ) + } + } else { + (lhs, rhs) + }; + if $unsigned { + _mm512_cmp_epu64_mask::(lhs, rhs) + } else { + _mm512_cmp_epi64_mask::(lhs, rhs) + } + }; + + // SAFETY: each output byte represents eight rows. The caller + // provides one complete output word for each 64 rows. + unsafe { output.add(chunk).write(mask ^ ($invert as u8)) }; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + } + }; +} + +impl_64!(i64, false, false); +impl_64!(u64, true, false); +impl_64!(f64, false, true); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_8.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_8.rs new file mode 100644 index 00000000000..832c10014e4 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/lanes_8.rs @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::*; + +use super::super::DenseComparison; +use super::super::DenseComparisonKind; +use super::super::SimdCompare; +use super::shared::ToSignedBits; + +const GREATER_THAN: i32 = _MM_CMPINT_NLE; + +macro_rules! impl_8 { + ($t:ty, $unsigned:expr) => { + #[allow(clippy::cast_possible_truncation)] + impl SimdCompare for $t { + const REQUIRES_AVX512BW: bool = true; + + #[target_feature(enable = "avx2")] + unsafe fn compare_words_avx2( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 2) { + let offset = chunk * 32; + let lhs = load_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi8, + __m256i + ); + let rhs = load_or_broadcast!( + $rhs, + offset, + $rhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi8, + __m256i + ); + let mask = if $eq { + _mm256_movemask_epi8(_mm256_cmpeq_epi8(lhs, rhs)) as u32 + } else { + let (lhs, rhs) = if $unsigned { + let sign = _mm256_set1_epi8(i8::MIN); + ( + _mm256_xor_si256(lhs, sign), + _mm256_xor_si256(rhs, sign), + ) + } else { + (lhs, rhs) + }; + _mm256_movemask_epi8(_mm256_cmpgt_epi8(lhs, rhs)) as u32 + }; + + // SAFETY: each output value represents 32 rows. The caller + // provides two values for each complete 64-row output word. + unsafe { output.add(chunk).write(mask ^ ($invert as u32)) }; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn compare_words_avx512( + comparison: DenseComparison<'_, Self>, + words: &mut [u64], + ) { + macro_rules! compare { + ( + lhs: $lhs:expr, + rhs: $rhs:expr, + lhs_constant: $lhs_constant:literal, + rhs_constant: $rhs_constant:literal, + equal: $eq:literal, + invert: $invert:expr, + words: $words:expr, + ) => {{ + for (word_index, word) in $words.iter_mut().enumerate() { + let base = word_index * 64; + let lhs = load_or_broadcast!( + $lhs, + base, + $lhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi8, + __m512i + ); + let rhs = load_or_broadcast!( + $rhs, + base, + $rhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi8, + __m512i + ); + let mask = if $eq { + _mm512_cmpeq_epi8_mask(lhs, rhs) + } else if $unsigned { + _mm512_cmp_epu8_mask::(lhs, rhs) + } else { + _mm512_cmp_epi8_mask::(lhs, rhs) + }; + *word = mask ^ $invert; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + } + }; +} + +impl_8!(i8, false); +impl_8!(u8, true); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/mod.rs new file mode 100644 index 00000000000..4eb38911ba7 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/mod.rs @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! x86-64 primitive comparison kernels grouped by physical lane width. +//! +//! Each lane module keeps a primitive type's AVX2 and AVX-512 implementations together because +//! both methods belong to one [`SimdCompare`] implementation. [`shared`] owns the load, broadcast, +//! and normalized-comparison macros used by every width. [`float_order`] owns the transformations +//! from IEEE bit patterns to signed integer keys with Vortex's total ordering. +//! Each kernel writes its mask instruction's natural chunk directly into LSB-first bitmap storage. +//! +//! [`SimdCompare`]: super::SimdCompare + +#[macro_use] +mod shared; + +mod float_order; + +mod lanes_16; +mod lanes_32; +mod lanes_64; +mod lanes_8; diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/shared.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/shared.rs new file mode 100644 index 00000000000..72ed77c34ee --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd/x86/shared.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use half::f16; + +pub(super) trait ToSignedBits { + fn to_signed_bits(self) -> i64; +} + +macro_rules! impl_integer_to_signed_bits { + ($($t:ty),+ $(,)?) => { + $( + impl ToSignedBits for $t { + fn to_signed_bits(self) -> i64 { + self as i64 + } + } + )+ + }; +} + +impl_integer_to_signed_bits!(i8, u8, i16, u16, i32, u32, i64, u64); + +impl ToSignedBits for f16 { + fn to_signed_bits(self) -> i64 { + i64::from(self.to_bits()) + } +} + +impl ToSignedBits for f32 { + fn to_signed_bits(self) -> i64 { + i64::from(self.to_bits()) + } +} + +impl ToSignedBits for f64 { + fn to_signed_bits(self) -> i64 { + self.to_bits() as i64 + } +} + +/// Pack the even bits of a 32-bit mask into the low 16 bits. +/// +/// AVX2 produces one sign bit for each byte. A 16-bit comparison repeats each lane result in two +/// adjacent bytes, so the masks successively merge groups of 1, 2, 4, and 8 even-positioned bits. +#[inline(always)] +pub(super) fn compact_even_bits(mut mask: u32) -> u64 { + mask &= 0x5555_5555; + mask = (mask | (mask >> 1)) & 0x3333_3333; + mask = (mask | (mask >> 2)) & 0x0f0f_0f0f; + mask = (mask | (mask >> 4)) & 0x00ff_00ff; + ((mask | (mask >> 8)) & 0x0000_ffff) as u64 +} + +macro_rules! load_or_broadcast { + ($input:expr, $base:expr, $constant:literal, $load:ident, $set:ident, $vec:ty) => {{ + if $constant { + // SAFETY: `dispatch_comparison` passes a pointer to a live local constant and consumes + // it within the same macro expansion. + $set(unsafe { *$input }.to_signed_bits() as _) + } else { + // SAFETY: each input covers `words.len() * 64` rows. The 8-, 16-, 32-, and 64-bit lane + // loops stop before the loaded vector's last lane reaches that row count. + unsafe { $load($input.add($base).cast::<$vec>()) } + } + }}; +} + +macro_rules! dispatch_comparison { + ($comparison:expr, $words:expr, $compare:ident) => { + match $comparison.kind { + DenseComparisonKind::EqualArrays { lhs, rhs, invert } => { + let invert = u64::from(invert).wrapping_neg(); + $compare!( + lhs: lhs.as_ptr(), + rhs: rhs.as_ptr(), + lhs_constant: false, + rhs_constant: false, + equal: true, + invert: invert, + words: $words, + ) + } + DenseComparisonKind::GreaterArrays { lhs, rhs, invert } => { + let invert = u64::from(invert).wrapping_neg(); + $compare!( + lhs: lhs.as_ptr(), + rhs: rhs.as_ptr(), + lhs_constant: false, + rhs_constant: false, + equal: false, + invert: invert, + words: $words, + ) + } + DenseComparisonKind::EqualArrayConstant { + values, + constant, + invert, + } => { + let invert = u64::from(invert).wrapping_neg(); + $compare!( + lhs: values.as_ptr(), + rhs: std::ptr::from_ref(&constant), + lhs_constant: false, + rhs_constant: true, + equal: true, + invert: invert, + words: $words, + ) + } + DenseComparisonKind::GreaterArrayConstant { + values, + constant, + invert, + } => { + let invert = u64::from(invert).wrapping_neg(); + $compare!( + lhs: values.as_ptr(), + rhs: std::ptr::from_ref(&constant), + lhs_constant: false, + rhs_constant: true, + equal: false, + invert: invert, + words: $words, + ) + } + DenseComparisonKind::GreaterConstantArray { + constant, + values, + invert, + } => { + let invert = u64::from(invert).wrapping_neg(); + $compare!( + lhs: std::ptr::from_ref(&constant), + rhs: values.as_ptr(), + lhs_constant: true, + rhs_constant: false, + equal: false, + invert: invert, + words: $words, + ) + } + } + }; +} + +#[cfg(test)] +mod tests { + use super::compact_even_bits; + + #[test] + fn test_compact_even_bits_matches_scalar_reference() { + for mask in [ + 0, // No bits are set. + u32::MAX, // All bits are set. + 0x5555_5555, // Every even bit is set. + 0xaaaa_aaaa, // Every odd bit is set. + 0xdead_beef, // An arbitrary mixture is set. + ] { + let expected = + (0..16).fold(0, |packed, bit| packed | (((mask >> (bit * 2)) & 1) << bit)); + + assert_eq!( + compact_even_bits(mask), + u64::from(expected), + "mask={mask:#x}" + ); + } + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/tests.rs new file mode 100644 index 00000000000..ed26779cbb1 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/tests.rs @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::NativePType; +use crate::dtype::half::f16; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::operators::Operator; + +fn apply_reference(lhs: T, rhs: T, op: Operator) -> bool { + match op { + Operator::Eq => lhs.is_eq(rhs), + Operator::NotEq => !lhs.is_eq(rhs), + Operator::Gt => lhs.is_gt(rhs), + Operator::Gte => lhs.is_ge(rhs), + Operator::Lt => lhs.is_lt(rhs), + Operator::Lte => lhs.is_le(rhs), + _ => unreachable!("primitive comparison test requires a comparison operator"), + } +} + +fn assert_primitive_comparison( + len: usize, + op: Operator, + lhs_pattern: [T; 4], + rhs_pattern: [T; 4], +) -> VortexResult<()> +where + T: NativePType, + Scalar: From, +{ + let mut ctx = array_session().create_execution_ctx(); + let lhs_values = lhs_pattern + .into_iter() + .cycle() + .take(len) + .collect::>(); + let rhs_values = rhs_pattern + .into_iter() + .cycle() + .take(len) + .collect::>(); + + let expected = lhs_values + .iter() + .copied() + .zip(rhs_values.iter().copied()) + .map(|(lhs, rhs)| apply_reference(lhs, rhs, op)) + .collect::>(); + let result = PrimitiveArray::from_iter(lhs_values.clone()) + .into_array() + .binary( + PrimitiveArray::from_iter(rhs_values.clone()).into_array(), + op, + )?; + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + + let constant = lhs_pattern[1]; + let expected = lhs_values + .iter() + .copied() + .map(|lhs| apply_reference(lhs, constant, op)) + .collect::>(); + let result = PrimitiveArray::from_iter(lhs_values) + .into_array() + .binary(ConstantArray::new(constant, len).into_array(), op)?; + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + + let expected = rhs_values + .iter() + .copied() + .map(|rhs| apply_reference(constant, rhs, op)) + .collect::>(); + let result = ConstantArray::new(constant, len) + .into_array() + .binary(PrimitiveArray::from_iter(rhs_values).into_array(), op)?; + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + +#[test] +fn test_compare_primitive_type_operator_and_bitmap_boundaries() -> VortexResult<()> { + const OPERATORS: [Operator; 6] = [ + Operator::Eq, + Operator::NotEq, + Operator::Gt, + Operator::Gte, + Operator::Lt, + Operator::Lte, + ]; + const LHS_FLOAT_PATTERN: [f32; 4] = [ + 1.0, // Less than the right value. + 5.0, // Equal to the right value. + 9.0, // Greater than the right value. + 2.0, // Less than the right value. + ]; + const RHS_FLOAT_PATTERN: [f32; 4] = [ + 3.0, // Greater than the left value. + 5.0, // Equal to the left value. + 7.0, // Less than the left value. + 4.0, // Greater than the left value. + ]; + + macro_rules! assert_integer_types { + ($len:expr, $op:expr, $($T:ty),+ $(,)?) => { + $(assert_primitive_comparison::<$T>( + $len, + $op, + [ + 1, // Less than the right value. + 5, // Equal to the right value. + 9, // Greater than the right value. + 2, // Less than the right value. + ], + [ + 3, // Greater than the left value. + 5, // Equal to the left value. + 7, // Less than the left value. + 4, // Greater than the left value. + ], + )?;)+ + }; + } + + for len in [ + 63, // One row before a complete bitmap word. + 64, // One complete bitmap word. + 65, // One row after a complete bitmap word. + ] { + for op in OPERATORS { + assert_integer_types!(len, op, u8, u16, u32, u64, i8, i16, i32, i64); + assert_primitive_comparison::( + len, + op, + LHS_FLOAT_PATTERN.map(f16::from_f32), + RHS_FLOAT_PATTERN.map(f16::from_f32), + )?; + assert_primitive_comparison::(len, op, LHS_FLOAT_PATTERN, RHS_FLOAT_PATTERN)?; + assert_primitive_comparison::( + len, + op, + LHS_FLOAT_PATTERN.map(f64::from), + RHS_FLOAT_PATTERN.map(f64::from), + )?; + } + } + + Ok(()) +} + +#[test] +fn test_compare_primitive_empty() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = PrimitiveArray::from_iter(Vec::::new()).into_array(); + let rhs = PrimitiveArray::from_iter(Vec::::new()).into_array(); + + let result = lhs.binary(rhs, Operator::Eq)?; + + assert_arrays_eq!(result, BoolArray::from_iter(Vec::::new()), &mut ctx); + + Ok(()) +} + +#[test] +fn test_compare_primitive_nullable_bitmap_boundary() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = (0..65) + .map(|index| (!matches!(index, 0 | 63)).then_some(index as i64)) + .collect::>(); + let rhs = (0..65) + .map(|index| (!matches!(index, 1 | 64)).then_some((64 - index) as i64)) + .collect::>(); + let expected = lhs + .iter() + .copied() + .zip(rhs.iter().copied()) + .map(|(lhs, rhs)| lhs.zip(rhs).map(|(lhs, rhs)| lhs < rhs)) + .collect::>(); + + let result = PrimitiveArray::from_option_iter(lhs).into_array().binary( + PrimitiveArray::from_option_iter(rhs).into_array(), + Operator::Lt, + )?; + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +}