From e4d3617995cf6b908a13d6315636968913b0bf0c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 22 Aug 2026 12:36:49 -0400 Subject: [PATCH] Use RowFn for primitive numeric comparisons Signed-off-by: Connor Tsui --- Cargo.lock | 1 + vortex-array/Cargo.toml | 5 + vortex-array/benches/compare.rs | 115 ++ vortex-array/benches/compare_primitive.rs | 482 ++++++ .../src/scalar_fn/fns/binary/compare/mod.rs | 23 +- .../scalar_fn/fns/binary/compare/primitive.rs | 180 +-- .../fns/binary/compare/primitive/columnar.rs | 121 ++ .../fns/binary/compare/primitive/simd.rs | 1398 +++++++++++++++++ .../src/scalar_fn/fns/binary/compare/tests.rs | 262 ++- vortex-array/src/scalar_fn/fns/binary/mod.rs | 7 +- vortex-array/src/test_harness/mod.rs | 24 + 11 files changed, 2515 insertions(+), 103 deletions(-) create mode 100644 vortex-array/benches/compare_primitive.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd.rs diff --git a/Cargo.lock b/Cargo.lock index 0010a4e6cec..2508ba7a3e8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9631,6 +9631,7 @@ dependencies = [ "uuid", "vortex-array", "vortex-array-macros", + "vortex-bench-support", "vortex-buffer", "vortex-compute", "vortex-error", diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 7c5936627e5..8df13ae0ad2 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -97,6 +97,7 @@ vortex-array = { path = ".", features = [ "table-display", "unstable_row_fns", ] } +vortex-bench-support = { workspace = true } [[bench]] name = "aggregate_max" @@ -130,6 +131,10 @@ harness = false name = "compare" harness = false +[[bench]] +name = "compare_primitive" +harness = false + [[bench]] name = "binary_ops" harness = false diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 43969a91e39..c1facb0adb9 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -4,6 +4,7 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; @@ -39,6 +40,7 @@ const ARRAY_SIZE: usize = 8_192; fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); bencher + .counter(ItemsCount::new(ARRAY_SIZE)) .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input @@ -50,6 +52,31 @@ fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { }); } +fn u8_array(offset: u8) -> ArrayRef { + (0u8..=u8::MAX) + .cycle() + .take(ARRAY_SIZE) + .map(|value| value.wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn i32_array(offset: i32) -> ArrayRef { + (0i32..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn u64_array(offset: u64) -> ArrayRef { + (0u64..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + fn bool_array(rng: &mut StdRng) -> ArrayRef { BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() } @@ -88,6 +115,13 @@ fn float_array(rng: &mut StdRng) -> ArrayRef { .into_array() } +fn f32_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f32..1.0)) + .collect::>() + .into_array() +} + fn string_array(rng: &mut StdRng) -> ArrayRef { VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { let len = rng.random_range(1usize..24); @@ -154,6 +188,14 @@ fn compare_int_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Gte); } +#[divan::bench] +fn compare_int_constant_lhs(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + let arr = int_array(&mut rng); + bench_compare(bencher, constant, arr, Operator::Gte); +} + #[divan::bench] fn compare_int_eq(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -162,6 +204,55 @@ fn compare_int_eq(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Eq); } +#[divan::bench] +fn compare_i32(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = i32_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_i32_constant(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = ConstantArray::new(1_000_000i32, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = u8_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8_constant(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = ConstantArray::new(127u8, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_constant(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = ConstantArray::new(1_000_000u64, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_eq(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Eq); +} + #[divan::bench] fn compare_float(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -170,6 +261,30 @@ fn compare_float(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[divan::bench] +fn compare_float_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_f32_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + #[divan::bench] fn compare_decimal(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); diff --git a/vortex-array/benches/compare_primitive.rs b/vortex-array/benches/compare_primitive.rs new file mode 100644 index 00000000000..baff6b004b1 --- /dev/null +++ b/vortex-array/benches/compare_primitive.rs @@ -0,0 +1,482 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A/B benchmarks for primitive comparison implementations. +//! +//! The local matrix is split into filterable groups. `throughput` covers every physical type, +//! comparison operator, and operand shape at a large row count. `scaling`, `validity`, and +//! `boundaries` isolate the dimensions that make a full Cartesian product unnecessarily large. +//! Filter a group with, for example, `cargo bench -p vortex-array --bench compare_primitive -- +//! scaling`. + +#![expect(clippy::unwrap_used)] + +use std::fmt; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::match_each_native_ptype; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_array::test_harness::compare_primitive_columnar; +use vortex_array::test_harness::compare_primitive_rows; +use vortex_array::validity::Validity as ArrayValidity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +const PTYPES: &[PType] = &[ + PType::U8, + PType::U16, + PType::U32, + PType::U64, + PType::I8, + PType::I16, + PType::I32, + PType::I64, + PType::F16, + PType::F32, + PType::F64, +]; +const OPERATORS: &[CompareOperator] = &[ + CompareOperator::Eq, + CompareOperator::NotEq, + CompareOperator::Gt, + CompareOperator::Gte, + CompareOperator::Lt, + CompareOperator::Lte, +]; +const SHAPES: &[OperandShape] = &[ + OperandShape::ArrayArray, + OperandShape::ArrayConstant, + OperandShape::ConstantArray, + OperandShape::ConstantConstant, +]; + +fn main() { + divan::main(); +} + +#[derive(Clone, Copy)] +struct ComparisonCase { + ptype: PType, + op: CompareOperator, + shape: OperandShape, + validity: ValidityCase, + row_count: usize, +} + +impl ComparisonCase { + const fn new( + ptype: PType, + op: CompareOperator, + shape: OperandShape, + validity: ValidityCase, + row_count: usize, + ) -> Self { + Self { + ptype, + op, + shape, + validity, + row_count, + } + } +} + +impl fmt::Display for ComparisonCase { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{}/{}/{}/{}/len={}", + self.ptype, + operator_name(self.op), + self.shape, + self.validity, + self.row_count + ) + } +} + +#[derive(Clone, Copy)] +enum OperandShape { + ArrayArray, + ArrayConstant, + ConstantArray, + ConstantConstant, +} + +impl fmt::Display for OperandShape { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::ArrayArray => "aa", + Self::ArrayConstant => "ac", + Self::ConstantArray => "ca", + Self::ConstantConstant => "cc", + }) + } +} + +#[derive(Clone, Copy)] +enum ValidityCase { + NonNullable, + NullableAllValid, + NullableLhs, + NullableRhs, + NullableBoth, + Masked, +} + +impl fmt::Display for ValidityCase { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::NonNullable => "nonnull", + Self::NullableAllValid => "nullable_all_valid", + Self::NullableLhs => "nullable_lhs", + Self::NullableRhs => "nullable_rhs", + Self::NullableBoth => "nullable_both", + Self::Masked => "masked", + }) + } +} + +#[derive(Clone, Copy)] +enum OperandValidity { + NonNullable, + NullableAllValid, + NullableSparse, + Masked, +} + +type CompareFn = + fn(&ArrayRef, &ArrayRef, CompareOperator, &mut ExecutionCtx) -> VortexResult; + +fn throughput_cases() -> Vec { + let mut cases = Vec::with_capacity(PTYPES.len() * OPERATORS.len() * SHAPES.len()); + for &ptype in PTYPES { + for &op in OPERATORS { + for &shape in SHAPES { + cases.push(ComparisonCase::new( + ptype, + op, + shape, + ValidityCase::NonNullable, + 65_536, + )); + } + } + } + cases +} + +fn scaling_cases() -> Vec { + const ROW_COUNTS: &[usize] = &[128, 1_024, 4_096, 8_192, 16_384, 65_536, 1_048_576]; + const SCALING_PTYPES: &[PType] = &[PType::U8, PType::I64, PType::F64]; + const SCALING_OPERATORS: &[CompareOperator] = &[CompareOperator::Eq, CompareOperator::Gte]; + + let mut cases = Vec::new(); + for &ptype in SCALING_PTYPES { + for &op in SCALING_OPERATORS { + for &shape in SHAPES { + for &row_count in ROW_COUNTS { + cases.push(ComparisonCase::new( + ptype, + op, + shape, + ValidityCase::NonNullable, + row_count, + )); + } + } + } + } + cases +} + +fn validity_cases() -> Vec { + const VALIDITY_PTYPES: &[PType] = &[PType::I64, PType::F64]; + const VALIDITY_OPERATORS: &[CompareOperator] = &[CompareOperator::Eq, CompareOperator::Gte]; + const VALIDITIES: &[ValidityCase] = &[ + ValidityCase::NullableAllValid, + ValidityCase::NullableLhs, + ValidityCase::NullableRhs, + ValidityCase::NullableBoth, + ValidityCase::Masked, + ]; + const ROW_COUNTS: &[usize] = &[128, 8_192, 65_536]; + + let mut cases = Vec::new(); + for &ptype in VALIDITY_PTYPES { + for &op in VALIDITY_OPERATORS { + for &shape in SHAPES { + for &validity in VALIDITIES { + for &row_count in ROW_COUNTS { + cases.push(ComparisonCase::new(ptype, op, shape, validity, row_count)); + } + } + } + } + } + cases +} + +fn boundary_cases() -> Vec { + const BOUNDARY_PTYPES: &[PType] = &[PType::U8, PType::I64, PType::F64]; + const BOUNDARY_OPERATORS: &[CompareOperator] = &[CompareOperator::Eq, CompareOperator::Gte]; + const ROW_COUNTS: &[usize] = &[0, 1, 63, 64, 65, 127, 128, 129]; + + let mut cases = Vec::new(); + for &ptype in BOUNDARY_PTYPES { + for &op in BOUNDARY_OPERATORS { + for &shape in SHAPES { + for &row_count in ROW_COUNTS { + cases.push(ComparisonCase::new( + ptype, + op, + shape, + ValidityCase::NonNullable, + row_count, + )); + } + } + } + } + cases +} + +fn cpu_feature_cases() -> Vec { + [ + (PType::I64, CompareOperator::Eq, OperandShape::ArrayArray), + (PType::I64, CompareOperator::Gte, OperandShape::ArrayArray), + ( + PType::I64, + CompareOperator::Gte, + OperandShape::ArrayConstant, + ), + ( + PType::I64, + CompareOperator::Gte, + OperandShape::ConstantArray, + ), + (PType::U64, CompareOperator::Gte, OperandShape::ArrayArray), + (PType::F64, CompareOperator::Gte, OperandShape::ArrayArray), + ] + .into_iter() + .map(|(ptype, op, shape)| { + ComparisonCase::new(ptype, op, shape, ValidityCase::NonNullable, 65_536) + }) + .collect() +} + +#[cfg(not(codspeed))] +mod throughput { + use super::*; + + #[divan::bench(args = throughput_cases())] + fn row(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_rows); + } + + #[divan::bench(args = throughput_cases())] + fn columnar(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_columnar); + } +} + +#[cfg(not(codspeed))] +mod scaling { + use super::*; + + #[divan::bench(args = scaling_cases())] + fn row(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_rows); + } + + #[divan::bench(args = scaling_cases())] + fn columnar(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_columnar); + } +} + +#[cfg(not(codspeed))] +mod validity { + use super::*; + + #[divan::bench(args = validity_cases())] + fn row(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_rows); + } + + #[divan::bench(args = validity_cases())] + fn columnar(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_columnar); + } +} + +#[cfg(not(codspeed))] +mod boundaries { + use super::*; + + #[divan::bench(args = boundary_cases())] + fn row(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_rows); + } + + #[divan::bench(args = boundary_cases())] + fn columnar(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_columnar); + } +} + +mod cpu_features { + use super::*; + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = cpu_feature_cases())] + fn row(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_rows); + } + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = cpu_feature_cases())] + fn columnar(bencher: Bencher, case: ComparisonCase) { + bench_case(bencher, case, compare_primitive_columnar); + } +} + +fn bench_case(bencher: Bencher, case: ComparisonCase, compare: CompareFn) { + let (lhs, rhs) = make_inputs(case).unwrap(); + let session = vortex_array::array_session(); + + bencher + .counter(ItemsCount::new(case.row_count)) + .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) + .bench_refs(|(lhs, rhs, ctx)| { + let result = compare(lhs, rhs, case.op, ctx).unwrap(); + + result.execute::(ctx) + }); +} + +fn make_inputs(case: ComparisonCase) -> VortexResult<(ArrayRef, ArrayRef)> { + match_each_native_ptype!(case.ptype, |T| { make_typed_inputs::(case) }) +} + +fn make_typed_inputs(case: ComparisonCase) -> VortexResult<(ArrayRef, ArrayRef)> +where + T: NativePType + Into + Into, +{ + let (lhs_validity, rhs_validity) = match case.validity { + ValidityCase::NonNullable => (OperandValidity::NonNullable, OperandValidity::NonNullable), + ValidityCase::NullableAllValid => ( + OperandValidity::NullableAllValid, + OperandValidity::NullableAllValid, + ), + ValidityCase::NullableLhs => ( + OperandValidity::NullableSparse, + OperandValidity::NonNullable, + ), + ValidityCase::NullableRhs => ( + OperandValidity::NonNullable, + OperandValidity::NullableSparse, + ), + ValidityCase::NullableBoth => ( + OperandValidity::NullableSparse, + OperandValidity::NullableSparse, + ), + ValidityCase::Masked => (OperandValidity::Masked, OperandValidity::Masked), + }; + + let lhs = match case.shape { + OperandShape::ArrayArray | OperandShape::ArrayConstant => { + make_array::(case.row_count, 1, lhs_validity)? + } + OperandShape::ConstantArray | OperandShape::ConstantConstant => { + make_constant::(case.row_count, lhs_validity)? + } + }; + let rhs = match case.shape { + OperandShape::ArrayArray | OperandShape::ConstantArray => { + make_array::(case.row_count, 17, rhs_validity)? + } + OperandShape::ArrayConstant | OperandShape::ConstantConstant => { + make_constant::(case.row_count, rhs_validity)? + } + }; + + Ok((lhs, rhs)) +} + +fn make_array(len: usize, offset: usize, validity: OperandValidity) -> VortexResult +where + T: NativePType, +{ + let values = (0..len) + .map(|index| { + let value = ((index.wrapping_mul(31) + offset) % 127) as i64; + T::from_i64(value).unwrap() + }) + .collect::>(); + + Ok(match validity { + OperandValidity::NonNullable => values.into_array(), + OperandValidity::NullableAllValid => { + PrimitiveArray::new(values, ArrayValidity::AllValid).into_array() + } + OperandValidity::NullableSparse => { + PrimitiveArray::new(values, sparse_validity(len, offset)).into_array() + } + OperandValidity::Masked => { + MaskedArray::try_new(values.into_array(), sparse_validity(len, offset))?.into_array() + } + }) +} + +fn make_constant(len: usize, validity: OperandValidity) -> VortexResult +where + T: NativePType + Into + Into, +{ + let value = T::from_i64(63).unwrap(); + + Ok(match validity { + OperandValidity::NonNullable => ConstantArray::new(value, len).into_array(), + OperandValidity::NullableAllValid => { + ConstantArray::new(Scalar::primitive(value, Nullability::Nullable), len).into_array() + } + OperandValidity::NullableSparse | OperandValidity::Masked => MaskedArray::try_new( + ConstantArray::new(value, len).into_array(), + sparse_validity(len, 3), + )? + .into_array(), + }) +} + +fn sparse_validity(len: usize, offset: usize) -> ArrayValidity { + ArrayValidity::from_iter((0..len).map(|index| !(index + offset).is_multiple_of(10))) +} + +const fn operator_name(op: CompareOperator) -> &'static str { + match op { + CompareOperator::Eq => "eq", + CompareOperator::NotEq => "not_eq", + CompareOperator::Gt => "gt", + CompareOperator::Gte => "gte", + CompareOperator::Lt => "lt", + CompareOperator::Lte => "lte", + } +} 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..24b6221f666 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -4,12 +4,12 @@ //! 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, row or fused lane +//! kernels for primitives, lane kernels for decimals, binary views for strings and 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. +//! Floating point values compare with Vortex's total ordering, including signed zero and ordered +//! NaN bit patterns. Equality is bitwise, matching [`Scalar`] comparison semantics. use std::cmp::Ordering; @@ -48,6 +48,10 @@ mod bytes; mod decimal; mod nested; mod primitive; +#[cfg(any(test, feature = "_test-harness"))] +pub(crate) use primitive::compare_primitive_columnar; +#[cfg(any(test, feature = "_test-harness"))] +pub(crate) use primitive::compare_primitive_rows; #[cfg(test)] mod tests; @@ -211,7 +215,7 @@ fn compare_arrays( ) .into_array()), DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx), - DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), + DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, ctx), DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => { @@ -282,8 +286,13 @@ 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); + + 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.rs index 1247358dce1..6581bd33a46 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,130 +1,124 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays via bit-packing lane kernels. +//! Primitive comparison execution through [`RowFn`] and fused lane kernels. +//! +//! Production uses the RowFn implementation on x86-64, where it has explicit SIMD kernels. Other +//! targets retain the fused columnar implementation until RowFn has a competitive target-specific +//! kernel. -use vortex_buffer::BitBuffer; +#[cfg(any(not(target_arch = "x86_64"), test, feature = "_test-harness"))] +mod columnar; +mod simd; + +use half::f16; use vortex_error::VortexResult; -use vortex_error::vortex_bail; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; use crate::dtype::DType; -use crate::dtype::NativePType; -use crate::dtype::Nullability; use crate::dtype::PType; -use crate::match_each_native_ptype; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::compare::collect_bits; -use crate::scalar_fn::fns::binary::compare::collect_zip_bits; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; /// Compare two primitive arrays of the same [`PType`]. /// -/// Floats compare with Vortex's total ordering: `NaN` is the largest value, `-0.0 < +0.0`, and -/// equality is bitwise. +/// Floats compare with Vortex's total ordering, including signed zero and ordered NaN bit +/// patterns. Equality is bitwise. pub(super) fn compare_primitive( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - match_each_native_ptype!(ptype, |T| { - compare_primitive_typed::(lhs, rhs, op, nullability, ctx) - }) + #[cfg(target_arch = "x86_64")] + { + compare_primitive_rows(lhs, rhs, op, ctx) + } + + #[cfg(not(target_arch = "x86_64"))] + { + columnar::compare_primitive(lhs, rhs, op, ctx) + } } -fn compare_primitive_typed( +/// Compare primitives through the retained columnar benchmark baseline. +#[cfg(any(test, feature = "_test-harness"))] +pub(crate) fn compare_primitive_columnar( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let len = lhs.len(); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - if lhs.len() != rhs.len() { - vortex_bail!( - "compare operator requires equal lengths, got {} and {}", - lhs.len(), - rhs.len() - ); - } + columnar::compare_primitive(lhs, rhs, op, ctx) +} - let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; - - let bits = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slices(lhs, rhs, op), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => compare_slice_constant(lhs, *rhs, op), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slice_constant(rhs, *lhs, op.swap()), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Unreachable through `execute_compare` (constant-constant is folded there), but - // cheap to answer anyway. - BitBuffer::full(apply_op(*lhs, *rhs, op), len) - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { - return Ok( - ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) - .into_array(), - ); - } - }; +pub(crate) fn compare_primitive_rows( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); - Ok(BoolArray::try_new(bits, validity)?.into_array()) + execute_rows(&PrimitiveCompare, &op, &args, ctx) } -#[inline(always)] -fn apply_op(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), +/// Internal row execution for primitive comparison operators. +#[derive(Clone)] +struct PrimitiveCompare; + +impl RowFn for PrimitiveCompare { + type Options = CompareOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + // `PrimitiveCompare` is a private implementation detail of `Binary`: it is never registered + // or serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. + ScalarFnVTable::id(&Binary) } -} -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { - // Dispatch the operator outside the lane loop so each instantiation vectorizes a single - // branch-free predicate. - match op { - CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = + PType::try_from(args.first().ok_or_else(|| { + vortex_err!("a comparison operator takes two operands, got none") + })?)?; + + match ptype { + PType::U8 => visit_compare_simd::(*op, visitor), + PType::U16 => visit_compare_simd::(*op, visitor), + PType::U32 => visit_compare_simd::(*op, visitor), + PType::I64 => visit_compare_simd::(*op, visitor), + PType::U64 => visit_compare_simd::(*op, visitor), + PType::I8 => visit_compare_simd::(*op, visitor), + PType::I16 => visit_compare_simd::(*op, visitor), + PType::I32 => visit_compare_simd::(*op, visitor), + PType::F16 => visit_compare_simd::(*op, visitor), + PType::F32 => visit_compare_simd::(*op, visitor), + PType::F64 => visit_compare_simd::(*op, visitor), + } } } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { - match op { - CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), - CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), - CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), - CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), - CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), - CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), - } +fn visit_compare_simd(op: CompareOperator, visitor: V) -> VortexResult +where + T: simd::SimdCompare, + V: RowVisitor, +{ + visitor.visit_kernel::<(T, T), _>(simd::PrimitiveComparisonKernel::new(op)) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs new file mode 100644 index 00000000000..4310b8e9feb --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -0,0 +1,121 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fused comparison and bit-packing for primitive lanes. +//! +//! Production uses this implementation on targets where the RowFn implementation does not have an +//! explicit SIMD kernel. Tests retain it on every target as the benchmark and semantic baseline. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare primitive operands with one fused comparison and bit-packing loop. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let ptype = PType::try_from(lhs.dtype())?; + match_each_native_ptype!(ptype, |T| { + compare_primitive_typed::(lhs, rhs, op, ctx) + }) +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + let bits = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slices(lhs, rhs, op), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => compare_slice_constant(lhs, *rhs, op), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slice_constant(rhs, *lhs, op.swap()), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => BitBuffer::full(apply_op(*lhs, *rhs, op), len), + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn apply_op(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 compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + } +} + +fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |lhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |lhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |lhs: T| lhs.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |lhs: T| lhs.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |lhs: T| lhs.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |lhs: T| lhs.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd.rs new file mode 100644 index 00000000000..1b976b4303d --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/simd.rs @@ -0,0 +1,1398 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Explicit SIMD primitive comparison kernels that write bitmap words directly. + +use std::marker::PhantomData; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::dtype::NativePType; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::unstable::row::ArgView; +use crate::scalar_fn::unstable::row::DenseRows; +use crate::scalar_fn::unstable::row::PackedBoolOutput; +use crate::scalar_fn::unstable::row::RowKernel; + +#[derive(Clone, Copy)] +pub(super) struct PrimitiveComparisonKernel { + op: CompareOperator, + marker: PhantomData, +} + +impl PrimitiveComparisonKernel { + pub(super) fn new(op: CompareOperator) -> Self { + Self { + op, + marker: PhantomData, + } + } +} + +impl RowKernel<(T, T)> for PrimitiveComparisonKernel { + type Element = bool; + type Output = PackedBoolOutput; + + fn eval(&self, (lhs, rhs): (T, T)) -> bool { + apply_op(lhs, rhs, self.op) + } + + fn collect_dense(&self, rows: DenseRows<'_, (T, T)>) -> VortexResult { + let row_count = rows.len(); + let lhs = PrimitiveInput::from_arg_view(rows.inputs().0.view(), row_count)?; + let rhs = PrimitiveInput::from_arg_view(rows.inputs().1.view(), row_count)?; + let mut output = PackedBoolOutput::zeroed(row_count); + + compare_into_words(lhs, rhs, self.op, row_count, output.words_mut()); + Ok(output) + } +} + +#[derive(Clone, Copy)] +pub(super) enum PrimitiveInput<'a, T> { + Slice(&'a [T]), + Constant(T), +} + +impl<'a, T: NativePType> PrimitiveInput<'a, T> { + fn from_arg_view(view: ArgView<'a, T>, row_count: usize) -> VortexResult { + match view { + ArgView::Column(values) => { + vortex_ensure!( + values.len() == row_count, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(Self::Slice(values)) + } + ArgView::Constant(value) => { + vortex_ensure!( + value.len() == 1, + "a decoded batch constant does not contain exactly one row", + ); + + Ok(Self::Constant(value[0])) + } + } + } + + /// 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. + 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, + } + } +} + +fn compare_into_words( + lhs: PrimitiveInput<'_, T>, + rhs: PrimitiveInput<'_, T>, + op: CompareOperator, + row_count: usize, + words: &mut [u64], +) { + if let (PrimitiveInput::Constant(lhs), PrimitiveInput::Constant(rhs)) = (lhs, rhs) { + let value = apply_op(lhs, rhs, op); + words.fill(if value { u64::MAX } else { 0 }); + if value + && let Some(last) = words.last_mut() + && !row_count.is_multiple_of(64) + { + *last = (1u64 << (row_count % 64)) - 1; + } + return; + } + + let full_words = row_count / 64; + + #[cfg(target_arch = "x86_64")] + if avx512_available::() { + let comparison = DenseComparison::new(lhs, rhs, op); + // SAFETY: `avx512_available` proves the required features. Each full word addresses exactly + // 64 input rows, or broadcasts a validated constant. + unsafe { T::compare_words_avx512(comparison, &mut words[..full_words]) }; + } else if avx2_available() { + let comparison = DenseComparison::new(lhs, rhs, op); + // SAFETY: `avx2_available` proves AVX2 support. Each full word addresses exactly 64 input + // rows, or broadcasts a validated constant. + unsafe { T::compare_words_avx2(comparison, &mut words[..full_words]) }; + } else { + compare_words_scalar(lhs, rhs, op, &mut words[..full_words]); + } + + #[cfg(not(target_arch = "x86_64"))] + compare_words_scalar(lhs, rhs, op, &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: see above. + let rhs = unsafe { rhs.get_unchecked(base + bit) }; + packed | ((apply_op(lhs, rhs, op) as u64) << bit) + }); + } +} + +#[cfg(target_arch = "x86_64")] +fn avx512_available() -> bool { + #[cfg(feature = "_test-harness")] + { + cfg!(target_feature = "avx512f") + && (!T::REQUIRES_AVX512BW || cfg!(target_feature = "avx512bw")) + } + + #[cfg(not(feature = "_test-harness"))] + { + std::arch::is_x86_feature_detected!("avx512f") + && (!T::REQUIRES_AVX512BW || std::arch::is_x86_feature_detected!("avx512bw")) + } +} + +#[cfg(target_arch = "x86_64")] +fn avx2_available() -> bool { + #[cfg(feature = "_test-harness")] + { + cfg!(target_feature = "avx2") + } + + #[cfg(not(feature = "_test-harness"))] + { + std::arch::is_x86_feature_detected!("avx2") + } +} + +#[cfg(target_arch = "x86_64")] +#[derive(Clone, Copy)] +pub(super) enum DenseComparison<'a, T> { + EqualArrays { + lhs: &'a [T], + rhs: &'a [T], + invert: u64, + }, + GreaterArrays { + lhs: &'a [T], + rhs: &'a [T], + invert: u64, + }, + EqualArrayConstant { + values: &'a [T], + constant: T, + invert: u64, + }, + GreaterArrayConstant { + values: &'a [T], + constant: T, + invert: u64, + }, + GreaterConstantArray { + constant: T, + values: &'a [T], + invert: u64, + }, +} + +#[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(), + ); + } + + match (lhs, rhs, op) { + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Eq) => { + Self::EqualArrays { + lhs, + rhs, + invert: 0, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::NotEq) => { + Self::EqualArrays { + lhs, + rhs, + invert: u64::MAX, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Gt) => { + Self::GreaterArrays { + lhs, + rhs, + invert: 0, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Gte) => { + Self::GreaterArrays { + lhs: rhs, + rhs: lhs, + invert: u64::MAX, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Lt) => { + Self::GreaterArrays { + lhs: rhs, + rhs: lhs, + invert: 0, + } + } + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs), CompareOperator::Lte) => { + Self::GreaterArrays { + lhs, + rhs, + invert: u64::MAX, + } + } + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Eq, + ) => Self::EqualArrayConstant { + values, + constant, + invert: 0, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::NotEq, + ) => Self::EqualArrayConstant { + values, + constant, + invert: u64::MAX, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Gt, + ) => Self::GreaterArrayConstant { + values, + constant, + invert: 0, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Gte, + ) => Self::GreaterConstantArray { + constant, + values, + invert: u64::MAX, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Lt, + ) => Self::GreaterConstantArray { + constant, + values, + invert: 0, + }, + ( + PrimitiveInput::Slice(values), + PrimitiveInput::Constant(constant), + CompareOperator::Lte, + ) => Self::GreaterArrayConstant { + values, + constant, + invert: u64::MAX, + }, + (PrimitiveInput::Constant(_), PrimitiveInput::Constant(_), _) => unreachable!(), + (PrimitiveInput::Constant(_), PrimitiveInput::Slice(_), _) => unreachable!(), + } + } +} + +#[cold] +fn compare_words_scalar( + lhs: PrimitiveInput<'_, T>, + rhs: PrimitiveInput<'_, T>, + op: CompareOperator, + words: &mut [u64], +) { + for (word_index, word) in words.iter_mut().enumerate() { + let base = word_index * 64; + *word = (0..64).fold(0, |packed, bit| { + // SAFETY: every complete word addresses 64 validated rows. + let lhs = unsafe { lhs.get_unchecked(base + bit) }; + // SAFETY: see above. + let rhs = unsafe { rhs.get_unchecked(base + bit) }; + packed | ((apply_op(lhs, rhs, op) as u64) << bit) + }); + } +} + +fn apply_op(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), + } +} + +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 these requirements can execute an unsupported + /// instruction or read outside an input allocation. + #[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 these requirements can execute an unsupported instruction or read outside an + /// input allocation. + #[cfg(target_arch = "x86_64")] + unsafe fn compare_words_avx512(comparison: DenseComparison<'_, Self>, words: &mut [u64]); +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use std::arch::x86_64::*; + + use half::f16; + + use super::DenseComparison; + use super::SimdCompare; + + #[inline(always)] + 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 { + $set(unsafe { *$input } as _) + } else { + // SAFETY: the enclosing word kernel proves this vector load is in bounds. + unsafe { $load($input.add($base).cast::<$vec>()) } + } + }}; + } + + trait ToI16Bits { + fn to_i16_bits(self) -> i16; + } + + impl ToI16Bits for i16 { + fn to_i16_bits(self) -> i16 { + self + } + } + + impl ToI16Bits for u16 { + fn to_i16_bits(self) -> i16 { + self as i16 + } + } + + impl ToI16Bits for f16 { + fn to_i16_bits(self) -> i16 { + self.to_bits() as i16 + } + } + + trait ToI32Bits { + fn to_i32_bits(self) -> i32; + } + + impl ToI32Bits for i32 { + fn to_i32_bits(self) -> i32 { + self + } + } + + impl ToI32Bits for u32 { + fn to_i32_bits(self) -> i32 { + self as i32 + } + } + + impl ToI32Bits for f32 { + fn to_i32_bits(self) -> i32 { + self.to_bits() as i32 + } + } + + macro_rules! load_16_or_broadcast { + ($input:expr, $base:expr, $constant:literal, $load:ident, $set:ident, $vec:ty) => {{ + if $constant { + $set(unsafe { *$input }.to_i16_bits()) + } else { + // SAFETY: the enclosing word kernel proves this vector load is in bounds. + unsafe { $load($input.add($base).cast::<$vec>()) } + } + }}; + } + + macro_rules! load_32_or_broadcast { + ($input:expr, $base:expr, $constant:literal, $load:ident, $set:ident, $vec:ty) => {{ + if $constant { + $set(unsafe { *$input }.to_i32_bits()) + } else { + // SAFETY: the enclosing word kernel proves this vector load is in bounds. + unsafe { $load($input.add($base).cast::<$vec>()) } + } + }}; + } + + trait ToI64Bits { + fn to_i64_bits(self) -> i64; + } + + impl ToI64Bits for u8 { + fn to_i64_bits(self) -> i64 { + i64::from(self) + } + } + + impl ToI64Bits for i64 { + fn to_i64_bits(self) -> i64 { + self + } + } + + impl ToI64Bits for u64 { + fn to_i64_bits(self) -> i64 { + self as i64 + } + } + + impl ToI64Bits for f64 { + fn to_i64_bits(self) -> i64 { + self.to_bits() as i64 + } + } + + macro_rules! load_64_or_broadcast { + ($input:expr, $base:expr, $constant:literal, $load:ident, $set:ident, $vec:ty) => {{ + if $constant { + $set(unsafe { *$input }.to_i64_bits()) + } else { + // SAFETY: the enclosing word kernel proves this vector load is in bounds. + unsafe { $load($input.add($base).cast::<$vec>()) } + } + }}; + } + + macro_rules! dispatch_comparison { + ($comparison:expr, $words:expr, $compare:ident) => { + match $comparison { + DenseComparison::EqualArrays { lhs, rhs, invert } => { + $compare!( + lhs.as_ptr(), + rhs.as_ptr(), + false, + false, + true, + invert, + $words + ) + } + DenseComparison::GreaterArrays { lhs, rhs, invert } => { + $compare!( + lhs.as_ptr(), + rhs.as_ptr(), + false, + false, + false, + invert, + $words + ) + } + DenseComparison::EqualArrayConstant { + values, + constant, + invert, + } => $compare!( + values.as_ptr(), + std::ptr::from_ref(&constant), + false, + true, + true, + invert, + $words + ), + DenseComparison::GreaterArrayConstant { + values, + constant, + invert, + } => $compare!( + values.as_ptr(), + std::ptr::from_ref(&constant), + false, + true, + false, + invert, + $words + ), + DenseComparison::GreaterConstantArray { + constant, + values, + invert, + } => $compare!( + std::ptr::from_ref(&constant), + values.as_ptr(), + true, + false, + false, + invert, + $words + ), + } + }; + } + + // The kernels below store each natural mask chunk directly into the output allocation. This + // avoids assembling a word with shifts and ORs, which LLVM can turn into a slower vectorized + // reduction. The subword order is the bitmap byte order because x86-64 is little-endian. + 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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $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::<6>(lhs, rhs) + } else { + _mm512_cmp_epi8_mask::<6>(lhs, rhs) + }; + *word = mask ^ $invert; + } + }}; + } + + dispatch_comparison!(comparison, words, compare); + } + } + }; + } + + 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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $words:expr) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 4) { + let offset = chunk * 16; + let lhs = load_16_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi16, + __m256i + ); + let rhs = load_16_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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $words:expr) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 2) { + let offset = chunk * 32; + let lhs = load_16_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi16, + __m512i + ); + let rhs = load_16_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::<6>(lhs, rhs) + } else { + _mm512_cmp_epi16_mask::<6>(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); + } + } + }; + } + + 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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $words:expr) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 8) { + let offset = chunk * 8; + let lhs = load_32_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi32, + __m256i + ); + let rhs = load_32_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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $words:expr) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 4) { + let offset = chunk * 16; + let lhs = load_32_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi32, + __m512i + ); + let rhs = load_32_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::<6>(lhs, rhs) + } else { + _mm512_cmp_epi32_mask::<6>(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); + } + } + }; + } + + 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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $words:expr) => {{ + macro_rules! compare_chunk { + ($offset:expr) => {{ + let lhs = load_64_or_broadcast!( + $lhs, + $offset, + $lhs_constant, + _mm256_loadu_si256, + _mm256_set1_epi64x, + __m256i + ); + let rhs = load_64_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: this method's 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:expr, $rhs:expr, $lhs_constant:literal, $rhs_constant:literal, $eq:literal, $invert:expr, $words:expr) => {{ + let output = $words.as_mut_ptr().cast::(); + for chunk in 0..($words.len() * 8) { + let offset = chunk * 8; + let lhs = load_64_or_broadcast!( + $lhs, + offset, + $lhs_constant, + _mm512_loadu_si512, + _mm512_set1_epi64, + __m512i + ); + let rhs = load_64_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: this method's 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::<6>(lhs, rhs) + } else { + _mm512_cmp_epi64_mask::<6>(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); + } + } + }; + } + + #[target_feature(enable = "avx2")] + #[inline] + unsafe fn float_key_16(bits: __m256i) -> __m256i { + let negative = _mm256_srai_epi16::<15>(bits); + _mm256_xor_si256(bits, _mm256_srli_epi16::<1>(negative)) + } + + #[target_feature(enable = "avx512f,avx512bw")] + #[inline] + 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)) + } + + #[target_feature(enable = "avx2")] + #[inline] + unsafe fn float_key_32(bits: __m256i) -> __m256i { + let negative = _mm256_srai_epi32::<31>(bits); + _mm256_xor_si256(bits, _mm256_srli_epi32::<1>(negative)) + } + + #[target_feature(enable = "avx512f")] + #[inline] + 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)) + } + + #[target_feature(enable = "avx2")] + #[inline] + 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)) + } + + #[target_feature(enable = "avx512f")] + #[inline] + 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)) + } + + impl_8!(i8, false); + impl_8!(u8, true); + impl_16!(i16, false, false); + impl_16!(u16, true, false); + impl_16!(f16, false, true); + impl_32!(i32, false, false); + impl_32!(u32, true, false); + impl_32!(f32, false, true); + impl_64!(i64, false, false); + impl_64!(u64, true, false); + impl_64!(f64, false, true); +} + +#[cfg(not(target_arch = "x86_64"))] +mod portable { + use half::f16; + + use super::SimdCompare; + + macro_rules! impl_simd_compare { + ($($t:ty),+ $(,)?) => { + $(impl SimdCompare for $t {})+ + }; + } + + impl_simd_compare!(i8, u8, i16, u16, f16, i32, u32, f32, i64, u64, f64); +} + +#[cfg(all(test, target_arch = "x86_64"))] +#[allow(clippy::cast_possible_truncation, clippy::tests_outside_test_module)] +mod tests { + use std::fmt::Debug; + + use half::f16; + + use super::DenseComparison; + use super::PrimitiveInput; + use super::SimdCompare; + use super::apply_op; + #[cfg(feature = "_test-harness")] + use super::avx2_available; + #[cfg(feature = "_test-harness")] + use super::avx512_available; + 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 benchmark_dispatch_uses_compiled_features() { + assert_eq!(avx2_available(), cfg!(target_feature = "avx2")); + assert_eq!(avx512_available::(), cfg!(target_feature = "avx512f")); + assert_eq!( + avx512_available::(), + cfg!(target_feature = "avx512f") && cfg!(target_feature = "avx512bw") + ); + } + + fn expected( + lhs: PrimitiveInput<'_, T>, + rhs: PrimitiveInput<'_, T>, + op: CompareOperator, + ) -> u64 { + (0..64).fold(0, |word, bit| { + // SAFETY: every test slice has exactly 64 elements. + let lhs = unsafe { lhs.get_unchecked(bit) }; + // SAFETY: see above. + let rhs = unsafe { rhs.get_unchecked(bit) }; + word | ((apply_op(lhs, rhs, op) as u64) << bit) + }) + } + + fn assert_simd_matches_scalar(lhs: &[T; 64], rhs: &[T; 64]) + where + T: SimdCompare + Debug, + { + let shapes = [ + (PrimitiveInput::Slice(lhs), PrimitiveInput::Slice(rhs)), + (PrimitiveInput::Slice(lhs), PrimitiveInput::Constant(rhs[7])), + ( + PrimitiveInput::Constant(lhs[11]), + PrimitiveInput::Slice(rhs), + ), + ]; + + for op in OPS { + for (lhs, rhs) in shapes { + let expected = expected(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 {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 {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, + { + for len in [0, 1, 7, 8, 9, 63, 64, 65, 129] { + let lhs_values = lhs_seed + .iter() + .copied() + .cycle() + .take(len) + .collect::>(); + let rhs_values = rhs_seed + .iter() + .copied() + .cycle() + .take(len) + .collect::>(); + let shapes = [ + ( + PrimitiveInput::Slice(lhs_values.as_slice()), + PrimitiveInput::Slice(rhs_values.as_slice()), + ), + ( + PrimitiveInput::Slice(lhs_values.as_slice()), + PrimitiveInput::Constant(rhs_seed[7]), + ), + ( + PrimitiveInput::Constant(lhs_seed[11]), + PrimitiveInput::Slice(rhs_values.as_slice()), + ), + ( + PrimitiveInput::Constant(lhs_seed[11]), + PrimitiveInput::Constant(rhs_seed[7]), + ), + ]; + + for op in OPS { + for (lhs, rhs) in shapes { + let mut actual = vec![u64::MAX; len.div_ceil(64)]; + compare_into_words(lhs, rhs, op, len, &mut actual); + + let mut expected = vec![0; len.div_ceil(64)]; + for index in 0..len { + // SAFETY: each slice input contains exactly `len` values. + let lhs = unsafe { lhs.get_unchecked(index) }; + // SAFETY: see above. + let rhs = unsafe { rhs.get_unchecked(index) }; + expected[index / 64] |= (apply_op(lhs, rhs, op) as u64) << (index % 64); + } + + assert_eq!(actual, expected, "word loop len={len} op={op:?}"); + } + } + } + } + + #[test] + fn 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 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 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 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 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 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 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 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 float_64_total_order_masks_match_scalar() { + const BITS: [u64; 16] = [ + 0xfff8_0000_0000_0001, + 0xfff0_0000_0000_0000, + 0xbff0_0000_0000_0000, + 0x8000_0000_0000_0001, + 0x8000_0000_0000_0000, + 0x0000_0000_0000_0000, + 0x0000_0000_0000_0001, + 0x3ff0_0000_0000_0000, + 0x7ff0_0000_0000_0000, + 0x7ff8_0000_0000_0000, + 0x7ff8_0000_0000_0001, + 0x7fff_ffff_ffff_ffff, + 0xfff8_0000_0000_0000, + 0x4000_0000_0000_0000, + 0xc000_0000_0000_0000, + 0x3fe0_0000_0000_0000, + ]; + 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 float_16_total_order_masks_match_scalar() { + const BITS: [u16; 16] = [ + 0xfe01, 0xfc00, 0xbc00, 0x8001, 0x8000, 0x0000, 0x0001, 0x3c00, 0x7c00, 0x7e00, 0x7e01, + 0x7fff, 0xfe00, 0x4000, 0xc000, 0x3800, + ]; + 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 float_32_total_order_masks_match_scalar() { + const BITS: [u32; 16] = [ + 0xffc0_0001, + 0xff80_0000, + 0xbf80_0000, + 0x8000_0001, + 0x8000_0000, + 0x0000_0000, + 0x0000_0001, + 0x3f80_0000, + 0x7f80_0000, + 0x7fc0_0000, + 0x7fc0_0001, + 0x7fff_ffff, + 0xffc0_0000, + 0x4000_0000, + 0xc000_0000, + 0x3f00_0000, + ]; + 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/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index 9831a963354..58934ee217b 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -10,6 +10,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -20,6 +21,7 @@ use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; +use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinArray; @@ -32,13 +34,17 @@ use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldName; use crate::dtype::FieldNames; +use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; +use crate::match_each_native_ptype; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::primitive::compare_primitive_columnar; +use crate::scalar_fn::fns::binary::compare::primitive::compare_primitive_rows; use crate::scalar_fn::fns::binary::scalar_cmp; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -406,8 +412,9 @@ fn int_constant_lhs_and_rhs(#[case] op: Operator, #[case] expected: [bool; 4]) { assert_arrays_eq!(swapped, BoolArray::from_iter(expected_swapped), &mut ctx); } -/// Floats compare with Vortex's total ordering: NaN is the largest value, equality is bitwise, -/// and -0.0 < +0.0. This matches `Scalar` comparison semantics. +/// Floats compare with Vortex's total ordering: negative NaNs sort below negative infinity, +/// positive NaNs sort above positive infinity, equality is bitwise, and -0.0 < +0.0. This matches +/// `Scalar` comparison semantics. #[test] fn float_total_order() { let mut ctx = array_session().create_execution_ctx(); @@ -429,6 +436,257 @@ fn float_total_order() { ); } +#[rstest] +#[case::row_eq(compare_primitive_rows, CompareOperator::Eq)] +#[case::row_not_eq(compare_primitive_rows, CompareOperator::NotEq)] +#[case::row_lt(compare_primitive_rows, CompareOperator::Lt)] +#[case::columnar_eq(compare_primitive_columnar, CompareOperator::Eq)] +#[case::columnar_not_eq(compare_primitive_columnar, CompareOperator::NotEq)] +#[case::columnar_lt(compare_primitive_columnar, CompareOperator::Lt)] +fn test_primitive_comparison_paths_preserve_semantics( + #[case] compare: fn( + &ArrayRef, + &ArrayRef, + CompareOperator, + &mut ExecutionCtx, + ) -> VortexResult, + #[case] op: CompareOperator, +) -> VortexResult<()> { + let lhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::NAN, // Null on the left. + -0.0, // Signed zero ordering. + 1.0, // A finite value below NaN. + f64::NAN, // Null on the right. + ], + Validity::from_iter([ + true, // + false, // + true, // + true, // + true, // + ]), + ) + .into_array(); + let rhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::INFINITY, // Null on the left. + 0.0, // Signed zero ordering. + f64::NAN, // A finite value below NaN. + 1.0, // Null on the right. + ], + Validity::from_iter([ + true, // + true, // + true, // + true, // + false, // + ]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let actual = compare(&lhs, &rhs, op, &mut ctx)?; + let expected = match op { + CompareOperator::Eq => [ + Some(true), // Equal NaNs. + None, // Null on the left. + Some(false), // Distinct signed zeroes. + Some(false), // A finite value and NaN. + None, // Null on the right. + ], + CompareOperator::NotEq | CompareOperator::Lt => [ + Some(false), // Equal NaNs. + None, // Null on the left. + Some(true), // Distinct signed zeroes. + Some(true), // A finite value and NaN. + None, // Null on the right. + ], + _ => unreachable!(), + }; + let expected = BoolArray::from_iter(expected); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); + + Ok(()) +} + +#[test] +fn row_primitive_comparison_packs_bit_boundaries() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + for len in [0, 1, 63, 64, 65, 127, 128, 129] { + let lhs_values = (0..len).map(|index| index as u64).collect::>(); + let rhs_values = (0..len).map(|index| (index ^ 1) as u64).collect::>(); + let lhs = PrimitiveArray::from_iter(lhs_values.iter().copied()).into_array(); + let rhs = PrimitiveArray::from_iter(rhs_values.iter().copied()).into_array(); + + let actual = compare_primitive_rows(&lhs, &rhs, CompareOperator::Eq, &mut ctx)?; + let expected = BoolArray::from_iter( + lhs_values + .iter() + .zip(&rhs_values) + .map(|(lhs, rhs)| lhs == rhs), + ); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + let lhs_i64 = + PrimitiveArray::from_iter(lhs_values.iter().map(|value| *value as i64)).into_array(); + let rhs_i64 = + PrimitiveArray::from_iter(rhs_values.iter().map(|value| *value as i64)).into_array(); + let actual = compare_primitive_rows(&lhs_i64, &rhs_i64, CompareOperator::Gte, &mut ctx)?; + let expected = BoolArray::from_iter( + lhs_values + .iter() + .zip(&rhs_values) + .map(|(lhs, rhs)| lhs >= rhs), + ); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + let constant = ConstantArray::new(31_u64, len).into_array(); + let actual = compare_primitive_rows(&lhs, &constant, CompareOperator::Gt, &mut ctx)?; + let expected = BoolArray::from_iter(lhs_values.iter().map(|value| *value > 31)); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + let actual = compare_primitive_rows(&constant, &lhs, CompareOperator::Lt, &mut ctx)?; + let expected = BoolArray::from_iter(lhs_values.iter().map(|value| 31 < *value)); + assert_arrays_eq!(&actual, &expected, &mut ctx); + } + + let len = 65; + let lhs_validity = Validity::from_iter((0..len).map(|index| index % 3 != 0)); + let rhs_validity = Validity::from_iter((0..len).map(|index| index % 5 != 0)); + let lhs = MaskedArray::try_new( + ConstantArray::new(30_u64, len).into_array(), + lhs_validity.clone(), + )? + .into_array(); + let rhs = MaskedArray::try_new( + ConstantArray::new(31_u64, len).into_array(), + rhs_validity.clone(), + )? + .into_array(); + + let actual = compare_primitive_rows(&lhs, &rhs, CompareOperator::Lt, &mut ctx)?; + let expected = BoolArray::new(BitBuffer::new_set(len), lhs_validity.and(rhs_validity)?); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + Ok(()) +} + +#[test] +fn row_primitive_comparison_supports_every_native_type() -> VortexResult<()> { + for ptype in [ + PType::U8, + PType::U16, + PType::U32, + PType::U64, + PType::I8, + PType::I16, + PType::I32, + PType::I64, + PType::F16, + PType::F32, + PType::F64, + ] { + match_each_native_ptype!(ptype, |T| { assert_row_primitive_comparison_type::()? }); + } + + Ok(()) +} + +fn assert_row_primitive_comparison_type() -> VortexResult<()> { + let lhs_values = (0..129) + .map(|index| T::from_i64((index * 31 % 127) as i64).vortex_expect("value must fit")) + .collect::>(); + let rhs_values = (0..129) + .map(|index| T::from_i64((index * 17 % 127) as i64).vortex_expect("value must fit")) + .collect::>(); + let lhs = PrimitiveArray::from_iter(lhs_values.iter().copied()).into_array(); + let rhs = PrimitiveArray::from_iter(rhs_values.iter().copied()).into_array(); + let mut ctx = array_session().create_execution_ctx(); + + for op in [ + CompareOperator::Eq, + CompareOperator::NotEq, + CompareOperator::Gt, + CompareOperator::Gte, + CompareOperator::Lt, + CompareOperator::Lte, + ] { + let actual = compare_primitive_rows(&lhs, &rhs, op, &mut ctx)?; + let expected = BoolArray::from_iter(lhs_values.iter().zip(&rhs_values).map( + |(lhs, rhs)| 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), + }, + )); + assert_arrays_eq!(&actual, &expected, &mut ctx); + } + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[rstest] +#[case::i64_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::i64_not_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::u64_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::u64_not_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::f64_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::f64_not_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +fn test_primitive_equality_uses_row_fn_for_supported_ptype( + #[case] lhs: ArrayRef, + #[case] rhs: ArrayRef, + #[case] op: CompareOperator, + #[case] expected: [bool; 3], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let actual = compare_primitive_rows(&lhs, &rhs, op, &mut ctx)?; + + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + #[rstest] #[case(Operator::Eq, [true, false, true, true])] #[case(Operator::Lt, [false, true, false, false])] diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index ed31cdf8d46..5414b5f26fe 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -42,10 +42,15 @@ pub(crate) use boolean::execute_boolean; pub use boolean::kleene_boolean_buffer_scalar; pub use boolean::kleene_boolean_buffers; mod compare; +#[cfg(any(test, feature = "_test-harness"))] +pub(crate) use compare::compare_primitive_columnar; +#[cfg(any(test, feature = "_test-harness"))] +pub(crate) use compare::compare_primitive_rows; pub use compare::*; mod numeric; pub(crate) use numeric::*; -mod primitive_operand; +#[cfg(any(not(target_arch = "x86_64"), test, feature = "_test-harness"))] +pub(crate) mod primitive_operand; use crate::scalar::NumericOperator; use crate::scalar::Scalar; diff --git a/vortex-array/src/test_harness/mod.rs b/vortex-array/src/test_harness/mod.rs index d3a6b829e80..a08a288c538 100644 --- a/vortex-array/src/test_harness/mod.rs +++ b/vortex-array/src/test_harness/mod.rs @@ -8,9 +8,13 @@ use goldenfile::differs::binary_diff; use itertools::Itertools; use vortex_error::VortexResult; +use crate::ArrayRef; use crate::ExecutionCtx; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; +use crate::scalar_fn::fns::binary::compare_primitive_columnar as compare_columnar; +use crate::scalar_fn::fns::binary::compare_primitive_rows as compare_rows; +use crate::scalar_fn::fns::operators::CompareOperator; #[cfg(not(codspeed))] pub mod trace; @@ -40,3 +44,23 @@ pub fn to_int_indices(indices_bits: BoolArray, ctx: &mut ExecutionCtx) -> Vortex .filter_map(|(idx, v)| (v && mask.value(idx)).then_some(idx as u64)) .collect_vec()) } + +/// Compare primitive arrays through the row-function implementation. +pub fn compare_primitive_rows( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + compare_rows(lhs, rhs, op, ctx) +} + +/// Compare primitive arrays through the fused columnar implementation. +pub fn compare_primitive_columnar( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + compare_columnar(lhs, rhs, op, ctx) +}