diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 6a07d03f50b..ccd440ff939 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -38,10 +38,60 @@ static SESSION: LazyLock = LazyLock::new(array_session); const LEN: usize = 32_768; +const ROWFN_MATRIX_CASES: &[(usize, RowFnShape)] = &[ + (128, RowFnShape::PerRowPerRow), + (128, RowFnShape::PerRowConstant), + (128, RowFnShape::ConstantPerRow), + (128, RowFnShape::PerRowNullableConstant), + (LEN, RowFnShape::PerRowPerRow), + (LEN, RowFnShape::PerRowConstant), + (LEN, RowFnShape::ConstantPerRow), + (LEN, RowFnShape::PerRowNullableConstant), +]; + +#[derive(Clone, Copy, Debug)] +enum RowFnShape { + PerRowPerRow, + PerRowConstant, + ConstantPerRow, + PerRowNullableConstant, +} + /// Decimal Mul and Div cost far more per lane than Add, so they run over a shorter array to keep /// the instrumented CodSpeed runs quick. const DECIMAL_MUL_DIV_LEN: usize = 8_192; +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_add(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Add); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_subtract(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Sub); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_multiply(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Mul); +} + +fn bench_rowfn_shape(bencher: Bencher, len: usize, shape: RowFnShape, operator: Operator) { + let per_row = + || PrimitiveArray::from_iter((0..len).map(|index| (index % 1_024) as i64 + 1)).into_array(); + let constant = || ConstantArray::new(17_i64, len).into_array(); + let nullable_constant = || ConstantArray::new(Some(17_i64), len).into_array(); + + let (lhs, rhs) = match shape { + RowFnShape::PerRowPerRow => (per_row(), per_row()), + RowFnShape::PerRowConstant => (per_row(), constant()), + RowFnShape::ConstantPerRow => (constant(), per_row()), + RowFnShape::PerRowNullableConstant => (per_row(), nullable_constant()), + }; + + bench_primitive(bencher, lhs, rhs, operator); +} + #[divan::bench] fn add_i64_nonnull(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); @@ -170,6 +220,14 @@ fn div_i64_nonnull(bencher: Bencher) { bench_primitive(bencher, lhs, rhs, Operator::Div); } +#[divan::bench] +fn div_i64_nullable(bencher: Bencher) { + let lhs = primitive_nullable(1_000_000, 7).into_array(); + let rhs = primitive_nullable(17, 5).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn sub_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..054846b7ef7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure evidence without +//! scanning the finished output. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,33 +15,18 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. +/// Keep this wrapper inlineable so captured constants can become loop invariants in the caller. +/// The lane kernels retain their own inlining decisions. #[inline] pub(super) fn checked_lanes( source: S, @@ -61,7 +48,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +59,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - mut apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6dc0de0fbea..a0c427b142e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,16 +4,19 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a +//! [`RowFn`](crate::scalar_fn::unstable::row::RowFn), which owns null handling, constants, and +//! validity for them; see [`row`]. Decimal keeps its own columnar implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; -#[cfg(test)] -mod tests; +mod row; use decimal::execute_numeric_decimal; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -81,3 +84,6 @@ fn build_empty_result( Ok(Canonical::empty(&result_dtype).into_array()) } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 8fd53d15216..7cbe88c7c8e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,73 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! Checked arithmetic for one primitive row. + +use std::ops::BitOrAssign; + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::validity::Validity; - -struct CheckedAdd; -struct CheckedSub; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedMul; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedDiv; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. - const ERROR: &'static str; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; +/// OR-reducible evidence that a row failed, with [`Default`] meaning success. +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - /// How this operation reports a failing lane. See [`Failure`]. - type Failure: Failure; +impl Failure for bool {} +impl Failure for u8 {} +impl Failure for u16 {} +impl Failure for u32 {} +impl Failure for u64 {} - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. - fn apply(lhs: T, rhs: T) -> (T, Self::Failure); +/// One arithmetic operator at one width, split into its value and failure evidence. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. + const ERROR: &'static str; - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); + /// How this operation reports a failing row. See [`Failure`]. + type Fail: Failure; - (failed == Self::Failure::default()).then_some(value) - } + /// The result of this operation, paired with evidence of whether the row failed. + fn apply(lhs: T, rhs: T) -> (T, Self::Fail); } impl CheckedPrimitiveOp for CheckedAdd { const ERROR: &'static str = "integer overflow in checked add"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.add_value(rhs), lhs.add_error(rhs)) } @@ -76,9 +55,9 @@ impl CheckedPrimitiveOp for CheckedAdd { impl CheckedPrimitiveOp for CheckedSub { const ERROR: &'static str = "integer overflow in checked sub"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.sub_value(rhs), lhs.sub_error(rhs)) } @@ -87,9 +66,9 @@ impl CheckedPrimitiveOp for CheckedSub { impl CheckedPrimitiveOp for CheckedMul { const ERROR: &'static str = "integer overflow in checked mul"; - type Failure = T::MulFailure; + type Fail = T::MulFailure; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, T::MulFailure) { (lhs.mul_value(rhs), lhs.mul_failure(rhs)) } @@ -97,16 +76,10 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { let failed = lhs.div_error(rhs); let value = if failed { @@ -116,151 +89,16 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } } -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) -} - -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. +/// Per-width arithmetic used to compute values and failure evidence. /// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// The add, subtract, and multiply value methods **must** be total over every stored lane value. +/// [`Self::div_value`] may assume that [`Self::div_error`] returned `false` for the same operands. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// This may be a word rather than `bool` when narrowing evidence would block vectorization. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -269,18 +107,15 @@ trait CheckedArithmetic: NativePType { fn sub_error(self, rhs: Self) -> bool; fn mul_value(self, rhs: Self) -> Self; fn mul_failure(self, rhs: Self) -> Self::MulFailure; + + /// Divide operands that [`Self::div_error`] accepted. fn div_value(self, rhs: Self) -> Self; + + /// Return whether [`Self::div_value`] would trap for these operands. fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } -/// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// Generate the shared integer operations from their failure predicates. macro_rules! impl_checked_integer { ( $ty:ty, @@ -291,67 +126,57 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self.wrapping_add(rhs) } - #[inline(always)] + #[inline] fn add_error(self, rhs: Self) -> bool { let ($add_lhs, $add_rhs) = (self, rhs); $add_error } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self.wrapping_sub(rhs) } - #[inline(always)] + #[inline] fn sub_error(self, rhs: Self) -> bool { let ($sub_lhs, $sub_rhs) = (self, rhs); $sub_error } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self.wrapping_mul(rhs) } - #[inline(always)] + #[inline] $(#[$mul_failure_attr])* fn mul_failure(self, rhs: Self) -> $mul_failure_ty { let ($mf_lhs, $mf_rhs) = (self, rhs); $mul_failure } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, rhs: Self) -> bool { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// Unsigned multiplication reports its discarded high half as failure evidence. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -364,12 +189,7 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// Signed widths use a range check or discarded high-half evidence. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -377,9 +197,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -389,13 +206,17 @@ macro_rules! impl_checked_signed { let kept = wide as $ty; let discarded = (wide >> <$ty>::BITS) as $ty; + // A product fits exactly when its discarded half is the sign extension of the kept + // half. XOR reduces that comparison to zero evidence for success and nonzero evidence + // for overflow without converting the wide product to a branch. + (discarded ^ (kept >> (<$ty>::BITS - 1))) as $failure }); }; ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -408,7 +229,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -417,54 +238,47 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self + rhs } - #[inline(always)] + #[inline] fn add_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self - rhs } - #[inline(always)] + #[inline] fn sub_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self * rhs } - #[inline(always)] + #[inline] fn mul_failure(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -484,30 +298,27 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; - /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change - /// pivots, and a spread of magnitudes that straddles the 64-bit split. + /// Values around zero, signed extrema, and 32- and 64-bit boundaries where the discarded + /// multiplication half or its sign extension changes. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, // Additive identity. + 1, // Smallest positive value. + -1, // All sign bits set. + 2, // Small positive power of two. + -2, // Small negative power of two. + 3, // Small non-power of two. + i64::MIN, // Minimum signed value. + i64::MIN + 1, // Minimum signed value's neighbor. + i64::MAX, // Maximum signed value. + i64::MAX - 1, // Maximum signed value's neighbor. + 1 << 31, // First positive value outside i32. + 1 << 32, // First value with bit 32 set. + 1 << 62, // Largest positive power of two in i64. + -(1 << 62), // Negative counterpart of the largest power of two. + 0x7FFF_FFFF, // Maximum i32 represented as i64. + -0x8000_0000, // Minimum i32 represented as i64. ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { let failed = lhs.mul_failure(rhs) != ::default(); @@ -522,14 +333,11 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..b6f565d8ff4 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive arithmetic execution through [`RowFn`]. +//! +//! `Binary` keeps its registered contract; [`NumericBinary`] is only an execution helper. Decimal +//! arithmetic remains on its existing columnar path. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedArithmetic; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +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::unstable::row::InitializedElement; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::UninitElementSink; +use crate::scalar_fn::unstable::row::execute_rows; + +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + execute_rows(&NumericBinary, &op, &args, ctx) +} + +/// Internal row execution for the primitive arithmetic operators. +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + // Fallibility is queried without input dtypes, so this conservatively covers integer widths. + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + // `NumericBinary` 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 dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = PType::try_from( + args.first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?, + )?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_div::(visitor), + } + }) + } +} + +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + visitor.visit_deferred::<(T, T), T, Op::Fail>( + |(lhs, rhs)| Op::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(numeric_error(Op::ERROR)); + } + + Ok(()) + }, + ) +} + +fn visit_div(visitor: V) -> VortexResult +where + T: CheckedArithmetic, + V: RowVisitor, +{ + if T::PTYPE.is_float() { + return visit_checked::(visitor); + } + + // Integer division is scalar and expensive, so deferring its cheap failure check preserves no + // vectorization. Check each divide immediately and stop at the first failure. + // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when + // they need to skip invalid rows. + visitor.visit_into::<(T, T), UninitElementSink, _>(|(lhs, rhs), output| { + let (value, failed) = CheckedDiv::apply(lhs, rhs); + if failed { + return Err(numeric_error(>::ERROR)); + } + + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }) +} + +/// Keep rich error construction out of row closures so the closures remain inlineable. +#[cold] +#[inline(never)] +fn numeric_error(message: &'static str) -> VortexError { + vortex_err!(InvalidArgument: "{message}") +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..3813c8612b3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -201,8 +201,7 @@ fn test_integer_array_array_errors_on_valid_lanes() { assert!(result.is_err()); } -/// Multiply two non-nullable lanes of `lhs` by two of `rhs`, expecting `Some(product)` where the -/// product fits and `None` where the checked kernel must report overflow. +/// Assert one checked multiplication through the complete array execution path. #[track_caller] fn assert_multiply(lhs: T, rhs: T, expected: Option) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -297,13 +296,11 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// An overflow late in the batch must still be reported, unless its row is null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000;