diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index a878fcb3c05..730db283cd4 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -196,6 +196,7 @@ where lower, nullability, kernel, + false, &predicate, ctx, ) diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index 68389b6eb2b..cc4c7393554 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -14,6 +14,7 @@ use fastlanes::BitPacking; use fastlanes::BitPackingCompare; use fastlanes::FastLanesComparable; +use num_traits::AsPrimitive; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -73,7 +74,7 @@ impl CompareKernel for BitPacked { /// Compare every value against the constant via the fused FastLanes `unpack_cmp` kernel. /// /// `NativePType::is_eq` / `is_lt` etc. provide total comparison (matching the primitive between -/// kernel's dispatch shape). `NotEq` has no direct method, so use `!is_eq`. +/// kernel's dispatch shape). fn compare_constant_typed( lhs: ArrayView<'_, BitPacked>, bit_width: u8, @@ -85,88 +86,89 @@ fn compare_constant_typed( where T: NativePType + BitPackedIter + + FastLanesComparable::Physical> + + AsPrimitive<::Physical>, + ::Physical: BitPacking + + NativePType + + BitPackingCompare + FastLanesComparable::Physical>, - ::Physical: BitPacking + NativePType + BitPackingCompare, { - match operator { - CompareOperator::Eq => compare_with( - lhs, - bit_width, - rhs, - operator, - nullability, - |a, b| a.is_eq(b), - ctx, - ), - CompareOperator::NotEq => compare_with( - lhs, - bit_width, - rhs, - operator, - nullability, - |a, b| !a.is_eq(b), - ctx, - ), - CompareOperator::Lt => compare_with( - lhs, - bit_width, - rhs, - operator, - nullability, - |a, b| a.is_lt(b), - ctx, - ), - CompareOperator::Lte => compare_with( + // `NotEq`, `Gte` and `Lte` invert the result of `Eq`, `Lt` and `Gt`, so each pair shares one + // set of fused kernels. + let (base, negate) = match operator { + CompareOperator::Eq => (BaseCompare::Eq, false), + CompareOperator::NotEq => (BaseCompare::Eq, true), + CompareOperator::Lt => (BaseCompare::Lt, false), + CompareOperator::Gte => (BaseCompare::Lt, true), + CompareOperator::Gt => (BaseCompare::Gt, false), + CompareOperator::Lte => (BaseCompare::Gt, true), + }; + let cmp = move |value: T, other: T| base.apply(value, other); + let exact = move |value: T| cmp(value, rhs) != negate; + + // The encoder rejects negative integers, so packed values are never negative. A negative + // constant then compares the same against every packed value as against zero. + if rhs.is_lt(T::zero()) { + return constant_result(lhs, exact(T::zero()), nullability, &exact, ctx); + } + if let Some(value) = decided_by_bit_width(bit_width, rhs, operator, cmp) { + return constant_result(lhs, value != negate, nullability, &exact, ctx); + } + + // Both sides are now non-negative, so the unsigned physical values compare the same as the + // values of `T`. The fused kernels take function items of the physical type, so signed and + // unsigned arrays of one width share them. + let rhs = rhs.as_(); + match base { + BaseCompare::Eq => stream_compare_fused::( lhs, bit_width, rhs, - operator, nullability, - |a, b| a.is_le(b), + ::Physical::is_eq, + negate, + &exact, ctx, ), - CompareOperator::Gt => compare_with( + BaseCompare::Lt => stream_compare_fused::( lhs, bit_width, rhs, - operator, nullability, - |a, b| a.is_gt(b), + ::Physical::is_lt, + negate, + &exact, ctx, ), - CompareOperator::Gte => compare_with( + BaseCompare::Gt => stream_compare_fused::( lhs, bit_width, rhs, - operator, nullability, - |a, b| a.is_ge(b), + ::Physical::is_gt, + negate, + &exact, ctx, ), } } -fn compare_with( - lhs: ArrayView<'_, BitPacked>, - bit_width: u8, - rhs: T, - operator: CompareOperator, - nullability: Nullability, - cmp: F, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType - + BitPackedIter - + FastLanesComparable::Physical>, - ::Physical: BitPacking + NativePType + BitPackingCompare, - F: Fn(T, T) -> bool + Copy, -{ - let exact = move |value| cmp(value, rhs); - if let Some(value) = decided_by_bit_width(bit_width, rhs, operator, cmp) { - return constant_result(lhs, value, nullability, &exact, ctx); +/// The comparison that the fused kernel runs. The other operators are its negations. +#[derive(Clone, Copy)] +enum BaseCompare { + Eq, + Lt, + Gt, +} + +impl BaseCompare { + fn apply(self, value: T, other: T) -> bool { + match self { + Self::Eq => value.is_eq(other), + Self::Lt => value.is_lt(other), + Self::Gt => value.is_gt(other), + } } - stream_compare_fused(lhs, bit_width, rhs, nullability, cmp, &exact, ctx) } /// The result for every packed value, if the bit width alone decides it. @@ -253,7 +255,14 @@ mod tests { let packed = BitPackedData::encode(&prim.clone().into_array(), bw, &mut ctx)?; let rhs_val = (cap.min(2048) / 2) as $T; let rhs = ConstantArray::new(rhs_val, prim.len()).into_array(); - for op in [CompareOperator::Eq, CompareOperator::Lt, CompareOperator::Gte] { + for op in [ + CompareOperator::Eq, + CompareOperator::NotEq, + CompareOperator::Lt, + CompareOperator::Lte, + CompareOperator::Gt, + CompareOperator::Gte, + ] { let got = ::compare( packed.as_view(), &rhs, op, &mut ctx, )? @@ -279,6 +288,44 @@ mod tests { sweep!(sweep_i32, i32, 1, 16, 31); sweep!(sweep_i64, i64, 1, 32, 63); + /// Signed arrays use the unsigned physical kernels. A negative constant compares the same + /// against every packed value, and at the widest packing the packed values still never set + /// the sign bit. + #[rstest] + #[case::negative_constant(7, -3)] + #[case::negative_constant_widest(31, -3)] + #[case::positive_constant_widest(31, 1_000_000)] + #[case::min_constant(7, i32::MIN)] + fn signed_compare_matches_primitive( + #[case] bit_width: u8, + #[case] rhs_value: i32, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..2048) + .map(|i| if i % 7 == 0 { i32::MAX - i } else { i % 100 }) + .collect(); + let prim = PrimitiveArray::from_iter(values); + let packed = BitPackedData::encode(&prim.clone().into_array(), bit_width, &mut ctx)?; + let rhs = ConstantArray::new(rhs_value, prim.len()).into_array(); + for op in [ + CompareOperator::Eq, + CompareOperator::NotEq, + CompareOperator::Lt, + CompareOperator::Lte, + CompareOperator::Gt, + CompareOperator::Gte, + ] { + let got = ::compare(packed.as_view(), &rhs, op, &mut ctx)? + .expect("bit-packed compare kernel must engage"); + let want = prim + .clone() + .into_array() + .binary(rhs.clone(), Operator::from(op))?; + assert_arrays_eq!(got, want, &mut SESSION.create_execution_ctx()); + } + Ok(()) + } + /// Inline-patch path: encode signed i32 values that exceed the bit-width range so they /// end up in `Patches`. The streaming kernel must splice the patches in before the /// predicate runs. diff --git a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs index b7ab40db942..059cf8167a1 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs @@ -59,17 +59,25 @@ const WORDS_PER_CHUNK: usize = CHUNK_SIZE / U64_BITS; /// Compare the unpacked values of a [`BitPackedArray`] against `rhs` using the fused FastLanes /// `unpack_cmp` kernel, producing a [`BoolArray`]. /// -/// `kernel(value, rhs)` sees only packed values, so it may assume `[0, 2^W - 1]`. +/// `kernel(value, rhs)` sees only packed values, so it may assume `[0, 2^W - 1]`. It compares values +/// as `V`, which can be the physical type of `T` so that types of one width share kernels. When +/// `negate` is set, the result of `kernel` is inverted, so one kernel serves an operator and its +/// negation. /// `predicate(value)` is exact and handles patches and width 0; it is `&dyn` since that path is /// cold. /// /// [`BitPackedArray`]: crate::BitPackedArray -pub(super) fn stream_compare_fused( +#[expect( + clippy::too_many_arguments, + reason = "typed compare execution parameters" +)] +pub(super) fn stream_compare_fused( array: ArrayView<'_, BitPacked>, bit_width: u8, - rhs: T, + rhs: V, nullability: Nullability, kernel: K, + negate: bool, predicate: &dyn Fn(T) -> bool, ctx: &mut ExecutionCtx, ) -> VortexResult @@ -78,7 +86,8 @@ where + BitPackedIter + FastLanesComparable::Physical>, ::Physical: BitPacking + NativePType + BitPackingCompare, - K: Fn(T, T) -> bool + Copy, + V: FastLanesComparable::Physical>, + K: Fn(V, V) -> bool + Copy, { let len = array.len(); let bit_width = bit_width as usize; @@ -113,11 +122,14 @@ where // kernel assigns every word in `transposed`, so its previous contents are irrelevant. unsafe { <::Physical as BitPackingCompare>::unchecked_unpack_cmp::< - T, + V, _, >(bit_width, packed_chunk, &mut lane_major, kernel, rhs); } transpose_bits::<::Physical>(&lane_major, out); + if negate { + out.iter_mut().for_each(|word| *word = !*word); + } }, )?; }