[EXPERIMENTAL] perf: make checked lane chunking opt-in

vortex-array/src/scalar_fn/fns/binary/numeric/row.rs

@@ -63,6 +63,10 @@ impl RowFn for NumericBinary {
         ScalarFnVTable::id(&Binary)
     }
 
+    #[allow(
+        clippy::cognitive_complexity,
+        reason = "the primitive-type macro repeats the same operator dispatch for each type"
+    )]
     fn dispatch<V: RowVisitor>(
         &self,
         op: &Self::Options,
@@ -74,11 +78,20 @@ impl RowFn for NumericBinary {
                 .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?,
         )?;
 
+        // Keep this choice inside the type macro. An outer operator/type match loses ARM i64/u64
+        // vectorization by changing CGU placement. See the settings at `visit_checked_chunked`.
         match_each_native_ptype!(ptype, |T| {
             match op {
                 NumericOperator::Add => visit_checked::<T, CheckedAdd, V>(visitor),
                 NumericOperator::Sub => visit_checked::<T, CheckedSub, V>(visitor),
-                NumericOperator::Mul => visit_checked::<T, CheckedMul, V>(visitor),
+                NumericOperator::Mul => {
+                    // Chunking improves these signed widths across the measured x86 and ARM targets.
+                    if matches!(T::PTYPE, PType::I16 | PType::I32) {
+                        visit_checked_chunked::<T, CheckedMul, V>(visitor)
+                    } else {
+                        visit_checked::<T, CheckedMul, V>(visitor)
+                    }
+                }
                 NumericOperator::Div => visit_div::<T, V>(visitor),
             }
         })
@@ -103,6 +116,27 @@ where
     )
 }
 
+// Keep the visit bodies separate. With rustc 1.98.0, LLVM 22.1.8, 16 CGUs, LTO off, and +neon,
+// sharing callback construction and dispatch here loses the two-lane dense i64/u64 loops.
+// Recheck binary_ops multiplication IR and timings before combining these helpers.
+fn visit_checked_chunked<T, Op, V>(visitor: V) -> VortexResult<V::VisitResult>
+where
+    T: NativePType,
+    Op: CheckedPrimitiveOp<T>,
+    V: RowVisitor,
+{
+    visitor.visit_deferred_chunked::<(T, T), T, Op::Fail>(
+        |(lhs, rhs)| Op::apply(lhs, rhs),
+        |failure| {
+            if failure != <Op::Fail as Default>::default() {
+                return Err(numeric_error(Op::ERROR));
+            }
+
+            Ok(())
+        },
+    )
+}
+
 fn visit_div<T, V>(visitor: V) -> VortexResult<V::VisitResult>
 where
     T: CheckedArithmetic,

vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs

@@ -296,30 +296,54 @@ fn test_multiply_overflow_on_null_lane_ignored<T: NativePType>(
     Ok(())
 }
 
-/// An overflow late in the batch must still be reported, unless its row is null.
+/// Failures at a chunk boundary or in the remainder must survive reduction unless the row is null.
 #[rstest]
-#[case::reported(true)]
-#[case::suppressed_by_null(false)]
-fn test_multiply_overflow_survives_batch_reduction(
-    #[case] lane_is_valid: bool,
+#[case::unsigned_32(u32::MAX, 3, 9)]
+#[case::signed_16(i16::MAX, 3, 9)]
+#[case::signed_32(i32::MAX, 3, 9)]
+fn test_multiply_overflow_survives_batch_reduction<T: NativePType + Into<Scalar>>(
+    #[case] overflowing: T,
+    #[case] rhs: T,
+    #[case] product: T,
+    #[values(false, true)] lane_is_valid: bool,
+    #[values(63, 64, 128)] overflow_at: usize,
+    #[values(false, true)] constant_rhs: bool,
 ) -> VortexResult<()> {
-    const LEN: u32 = 1000;
-    const OVERFLOW_AT: u32 = 700;
+    const LEN: usize = 129;
 
     let mut ctx = array_session().create_execution_ctx();
-    let mut lhs: Vec<u32> = (0..LEN).map(|i| i % 100 + 1).collect();
-    lhs[OVERFLOW_AT as usize] = u32::MAX;
-
-    let validity = Validity::from_iter((0..LEN).map(|i| i != OVERFLOW_AT || lane_is_valid));
+    let mut lhs = vec![rhs; LEN];
+    lhs[overflow_at] = overflowing;
+    let validity = if lane_is_valid {
+        Validity::NonNullable
+    } else {
+        Validity::from_iter((0..LEN).map(|index| index != overflow_at))
+    };
+    let rhs = if constant_rhs {
+        ConstantArray::new(rhs, LEN).into_array()
+    } else {
+        PrimitiveArray::from_iter(vec![rhs; LEN]).into_array()
+    };
     let result = PrimitiveArray::new(Buffer::copy_from(&lhs), validity)
         .into_array()
-        .binary(
-            PrimitiveArray::from_iter(vec![3u32; LEN as usize]).into_array(),
-            Operator::Mul,
-        )?
+        .binary(rhs, Operator::Mul)?
         .execute::<RecursiveCanonical>(&mut ctx);
 
-    assert_eq!(result.is_err(), lane_is_valid);
+    if lane_is_valid {
+        assert!(
+            result.is_err(),
+            "overflow at row {overflow_at} must be reported"
+        );
+        return Ok(());
+    }
+
+    assert_arrays_eq!(
+        result?.0.into_array(),
+        PrimitiveArray::from_option_iter(
+            (0..LEN).map(|index| (index != overflow_at).then_some(product)),
+        ),
+        &mut ctx
+    );
 
     Ok(())
 }

vortex-array/src/scalar_fn/unstable/row/execute/mod.rs

@@ -8,6 +8,7 @@
 
 mod owned;
 pub(super) use owned::execute_owned;
+pub(super) use owned::execute_owned_chunked;
 pub(super) use owned::execute_owned_filtered;
 pub(super) use owned::execute_owned_infallible;
 pub(super) use owned::execute_owned_infallible_filtered;
@@ -22,6 +23,7 @@ pub(super) use packed_bool::execute_owned_infallible_bool;
 mod retry;
 pub(super) use retry::DenseAttempt;
 pub(super) use retry::execute_owned_dense_attempt;
+pub(super) use retry::execute_owned_dense_attempt_chunked;
 
 mod sink;
 pub(super) use sink::execute_sink;

vortex-array/src/scalar_fn/unstable/row/execute/owned.rs

@@ -300,3 +300,41 @@ where
     // SAFETY: normal completion of `map_checked_into` initializes every output slot.
     Ok(unsafe { values.finish(row_count, ctx.allocator()) })
 }
+
+/// Decode every input column, then reduce failure evidence within fixed-size chunks.
+///
+/// This stays separate from [`execute_owned`] to preserve the default loop's codegen context.
+pub(crate) fn execute_owned_chunked<Args, Out, Prepared, Fail>(
+    args: &dyn ExecutionArgs,
+    ctx: &mut ExecutionCtx,
+    prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
+    apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail),
+    finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
+) -> VortexResult<ArrayRef>
+where
+    Args: IndexedElementTuple,
+    Out: OutputElement,
+    Fail: FailureEvidence,
+{
+    // Errors and unwinds abandon partially initialized slots. The assertion ensures that no
+    // initialized value requires a destructor to run.
+    const { assert_owned_output_needs_no_drop::<Out>() };
+
+    let columns = Args::decode(args, ctx)?;
+    let prepared = prepare(Args::const_values(&columns));
+
+    let row_count = args.row_count();
+    let mut values = Out::with_capacity(row_count, ctx.allocator());
+    let output = &mut values.slots()[..row_count];
+
+    let Some(source) = decoded_source::<Args>(&columns, row_count) else {
+        vortex_bail!("a decoded row input does not address exactly {row_count} rows");
+    };
+    let failure = source.map_checked_chunked_into(output, |elements| apply(&prepared, elements));
+
+    // Defer rich error construction until after the row loop.
+    finish_failure(failure)?;
+
+    // SAFETY: normal completion of `map_checked_chunked_into` initializes every output slot.
+    Ok(unsafe { values.finish(row_count, ctx.allocator()) })
+}

vortex-array/src/scalar_fn/unstable/row/execute/retry.rs

@@ -111,3 +111,83 @@ where
         Err(error) => Ok(DenseAttempt::DeferredError(error)),
     }
 }
+
+/// Attempt dense execution with chunked failure reduction for non-constant inputs.
+///
+/// This follows the error contract of [`execute_owned_dense_attempt`].
+pub(in crate::scalar_fn::unstable::row) fn execute_owned_dense_attempt_chunked<
+    Args,
+    Out,
+    Prepared,
+    Fail,
+>(
+    args: &dyn ExecutionArgs,
+    ctx: &mut ExecutionCtx,
+    prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
+    apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail),
+    finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
+) -> VortexResult<DenseAttempt>
+where
+    Args: IndexedElementTuple,
+    Out: OutputElement,
+    Fail: FailureEvidence,
+{
+    // Errors and unwinds abandon partially initialized slots. The assertion ensures that no
+    // initialized value requires a destructor to run.
+    const { assert_owned_output_needs_no_drop::<Out>() };
+
+    // Keep this dense row loop separate from `execute_owned_chunked`. Factoring shared state into a
+    // helper changes LLVM's optimized dense kernel even when the helper is inlined.
+    let columns = Args::decode(args, ctx)?;
+    let prepared = prepare(Args::const_values(&columns));
+
+    let row_count = args.row_count();
+    let mut values = Out::with_capacity(row_count, ctx.allocator());
+    let output = &mut values.slots()[..row_count];
+
+    let failure_evidence = if let Some(views) = Args::views_if_no_consts(&columns) {
+        // Keep this validation beside the views so LLVM sees their common length here.
+        vortex_ensure!(
+            Args::view_lens_match(&views, row_count),
+            "a decoded row input does not address exactly {row_count} rows",
+        );
+
+        // SAFETY: `view_lens_match` checked that these exact retained views address `row_count`
+        // rows.
+        let source = unsafe { Args::indexed_source(views, row_count) };
+
+        source.map_checked_chunked_into(output, |elements| apply(&prepared, elements))
+    } else {
+        // Keep this proof branch-local. Shared validation prevents LLVM from specializing this
+        // loop for each batch-constant arrangement, leaving it scalar under multiple CGUs without
+        // LTO.
+        vortex_ensure!(
+            Args::decoded_lens_match(&columns, row_count),
+            "a decoded row input does not address exactly {row_count} rows",
+        );
+
+        let mut accumulated_failure = Fail::default();
+
+        // Iterate over `output` directly. A `0..row_count` range reuses the address-taken value
+        // from the validation error formatter and retains an output bounds check.
+        for (index, slot) in output.iter_mut().enumerate() {
+            // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the loop.
+            let (value, row_failure) = apply(&prepared, Args::get(&columns, index));
+
+            slot.write(value);
+            accumulated_failure |= row_failure;
+        }
+
+        accumulated_failure
+    };
+
+    match finish_failure(failure_evidence) {
+        Ok(()) => {
+            // SAFETY: normal completion of either path initializes every output slot.
+            let output = unsafe { values.finish(row_count, ctx.allocator()) };
+
+            Ok(DenseAttempt::Values(output))
+        }
+        Err(error) => Ok(DenseAttempt::DeferredError(error)),
+    }
+}

vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs

@@ -35,6 +35,7 @@ use crate::scalar_fn::unstable::row::RowFn;
 use crate::scalar_fn::unstable::row::SinkResult;
 use crate::scalar_fn::unstable::row::execute::execute_owned;
 use crate::scalar_fn::unstable::row::execute::execute_owned_bool;
+use crate::scalar_fn::unstable::row::execute::execute_owned_chunked;
 use crate::scalar_fn::unstable::row::execute::execute_owned_filtered;
 use crate::scalar_fn::unstable::row::execute::execute_owned_infallible;
 use crate::scalar_fn::unstable::row::execute::execute_owned_infallible_bool;
@@ -184,6 +185,34 @@ impl<F: RowFn> RowVisitor for ExecuteRows<'_, '_, F> {
         )
     }
 
+    fn visit_prepared_deferred_chunked<Args, Out, Prepared, Fail>(
+        self,
+        prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
+        apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail),
+        finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
+    ) -> VortexResult<Self::VisitResult>
+    where
+        Args: IndexedElementTuple,
+        Out: OutputElement,
+        Fail: FailureEvidence,
+    {
+        const { assert_deferred_visit_contract::<F, Args, Out, Fail>() };
+        let visited = BatchPlan::new(
+            validate_owned_visit::<Args, Out>(self.dtypes)?,
+            self.output_dtype,
+            RowPolicy::for_deferred_output::<Args>(),
+        )?;
+        self.plan.ensure_reproduced_by(&visited)?;
+
+        execute_owned_chunked::<Args, Out, Prepared, Fail>(
+            self.args,
+            self.ctx,
+            prepare,
+            apply,
+            finish_failure,
+        )
+    }
+
     fn visit_prepared_deferred_bool<Args, Prepared, Fail, const MULTIVERSIONED: bool>(
         self,
         prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,

vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs

@@ -33,6 +33,7 @@ use crate::scalar_fn::unstable::row::batch::BorrowedRowFnArgs;
 use crate::scalar_fn::unstable::row::execute::DenseAttempt;
 use crate::scalar_fn::unstable::row::execute::execute_bool_dense_attempt;
 use crate::scalar_fn::unstable::row::execute::execute_owned_dense_attempt;
+use crate::scalar_fn::unstable::row::execute::execute_owned_dense_attempt_chunked;
 use crate::scalar_fn::unstable::row::execute::execute_owned_infallible;
 use crate::scalar_fn::unstable::row::execute::execute_sink;
 
@@ -154,6 +155,34 @@ impl<F: RowFn> RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> {
         )
     }
 
+    fn visit_prepared_deferred_chunked<Args, Out, Prepared, Fail>(
+        self,
+        prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
+        apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail),
+        finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
+    ) -> VortexResult<Self::VisitResult>
+    where
+        Args: IndexedElementTuple,
+        Out: OutputElement,
+        Fail: FailureEvidence,
+    {
+        const { assert_deferred_visit_contract::<F, Args, Out, Fail>() };
+        let visited = BatchPlan::new(
+            validate_owned_visit::<Args, Out>(self.args.dtypes())?,
+            self.output_dtype,
+            RowPolicy::for_deferred_output::<Args>(),
+        )?;
+        self.args.plan().ensure_reproduced_by(&visited)?;
+
+        execute_owned_dense_attempt_chunked::<Args, Out, Prepared, Fail>(
+            self.args,
+            self.ctx,
+            prepare,
+            apply,
+            finish_failure,
+        )
+    }
+
     fn visit_prepared_deferred_bool<Args, Prepared, Fail, const MULTIVERSIONED: bool>(
         self,
         prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,