Skip to content

Commit a97717c

Browse files
joseph-isaacsclaude
andcommitted
Revert "refactor: collapse per-shape checked kernels into one care-lane dispatch"
This reverts commit 32c5287897299552896518135705421b9de85df9, restoring the per-shape checked kernels. Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv
1 parent 716ac07 commit a97717c

1 file changed

Lines changed: 179 additions & 44 deletions

File tree

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

Lines changed: 179 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -174,24 +174,15 @@ where
174174
(
175175
PrimitiveOperand::Array { values: lhs, .. },
176176
PrimitiveOperand::Array { values: rhs, .. },
177-
) => {
178-
let (lhs, rhs) = (lhs.as_slice(), rhs.as_slice());
179-
checked_op::<T, Op, _>(len, &care_lanes, |idx| (lhs[idx], rhs[idx]))
180-
}
177+
) => checked_array_array::<T, Op>(lhs, rhs, &care_lanes),
181178
(
182179
PrimitiveOperand::Array { values: lhs, .. },
183180
PrimitiveOperand::Constant { value: rhs, .. },
184-
) => {
185-
let (lhs, rhs) = (lhs.as_slice(), *rhs);
186-
checked_op::<T, Op, _>(len, &care_lanes, |idx| (lhs[idx], rhs))
187-
}
181+
) => checked_array_constant::<T, Op>(lhs, *rhs, &care_lanes),
188182
(
189183
PrimitiveOperand::Constant { value: lhs, .. },
190184
PrimitiveOperand::Array { values: rhs, .. },
191-
) => {
192-
let (lhs, rhs) = (*lhs, rhs.as_slice());
193-
checked_op::<T, Op, _>(len, &care_lanes, |idx| (lhs, rhs[idx]))
194-
}
185+
) => checked_constant_array::<T, Op>(*lhs, rhs, &care_lanes),
195186
(
196187
PrimitiveOperand::Constant { value: lhs, .. },
197188
PrimitiveOperand::Constant { value: rhs, .. },
@@ -311,48 +302,132 @@ impl<T: NativePType> CheckedValues<T> {
311302
}
312303
}
313304

314-
/// Dispatch a checked op over the care lanes, reading operands via `operands_at`.
315-
///
316-
/// The one-pass strategy short-circuits on the first observable error; the default
317-
/// split strategy computes every lane branch-free and re-scans only care lanes for
318-
/// errors, so undemanded and null lanes never fail the operation.
319-
fn checked_op<T, Op, F>(len: usize, care_lanes: &Mask, operands_at: F) -> CheckedValues<T>
305+
fn checked_array_array<T, Op>(lhs: &[T], rhs: &[T], care_lanes: &Mask) -> CheckedValues<T>
320306
where
321307
T: NativePType,
322308
Op: CheckedPrimitiveOp<T>,
323-
F: Fn(usize) -> (T, T) + Copy,
324309
{
310+
debug_assert_eq!(lhs.len(), rhs.len());
311+
325312
match care_lanes.bit_buffer() {
326-
AllOr::All if Op::CHECKED_VALUE_LOOP => checked_all_lanes(len, |idx| {
327-
let (lhs, rhs) = operands_at(idx);
328-
Op::checked(lhs, rhs)
329-
}),
330-
AllOr::All => collect_all_lanes(len, |idx| {
331-
let (lhs, rhs) = operands_at(idx);
332-
Op::apply(lhs, rhs)
333-
}),
334-
AllOr::None => CheckedValues::zeroed(len),
335-
AllOr::Some(care_bits) if Op::CHECKED_VALUE_LOOP => {
336-
checked_valid_lanes(len, care_bits, |idx| {
337-
let (lhs, rhs) = operands_at(idx);
338-
Op::checked(lhs, rhs)
339-
})
313+
AllOr::All if Op::CHECKED_VALUE_LOOP => checked_array_array_one_pass::<T, Op>(lhs, rhs),
314+
AllOr::All => checked_array_array_all_lanes::<T, Op>(lhs, rhs),
315+
AllOr::None => CheckedValues::zeroed(lhs.len()),
316+
AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => {
317+
checked_array_array_valid_lanes_one_pass::<T, Op>(lhs, rhs, valid_bits)
318+
}
319+
AllOr::Some(valid_bits) => checked_array_array_valid_lanes::<T, Op>(lhs, rhs, valid_bits),
320+
}
321+
}
322+
323+
fn checked_array_constant<T, Op>(lhs: &[T], rhs: T, care_lanes: &Mask) -> CheckedValues<T>
324+
where
325+
T: NativePType,
326+
Op: CheckedPrimitiveOp<T>,
327+
{
328+
match care_lanes.bit_buffer() {
329+
AllOr::All if Op::CHECKED_VALUE_LOOP => checked_array_constant_one_pass::<T, Op>(lhs, rhs),
330+
AllOr::All => checked_array_constant_all_lanes::<T, Op>(lhs, rhs),
331+
AllOr::None => CheckedValues::zeroed(lhs.len()),
332+
AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => {
333+
checked_array_constant_valid_lanes_one_pass::<T, Op>(lhs, rhs, valid_bits)
334+
}
335+
AllOr::Some(valid_bits) => {
336+
checked_array_constant_valid_lanes::<T, Op>(lhs, rhs, valid_bits)
337+
}
338+
}
339+
}
340+
341+
fn checked_constant_array<T, Op>(lhs: T, rhs: &[T], care_lanes: &Mask) -> CheckedValues<T>
342+
where
343+
T: NativePType,
344+
Op: CheckedPrimitiveOp<T>,
345+
{
346+
match care_lanes.bit_buffer() {
347+
AllOr::All if Op::CHECKED_VALUE_LOOP => checked_constant_array_one_pass::<T, Op>(lhs, rhs),
348+
AllOr::All => checked_constant_array_all_lanes::<T, Op>(lhs, rhs),
349+
AllOr::None => CheckedValues::zeroed(rhs.len()),
350+
AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => {
351+
checked_constant_array_valid_lanes_one_pass::<T, Op>(lhs, rhs, valid_bits)
340352
}
341-
AllOr::Some(care_bits) => {
342-
let mut checked = collect_all_lanes(len, |idx| {
343-
let (lhs, rhs) = operands_at(idx);
344-
Op::apply(lhs, rhs)
345-
});
346-
checked.failed = checked.failed
347-
&& any_valid_error(len, care_bits, |idx| {
348-
let (lhs, rhs) = operands_at(idx);
349-
Op::apply(lhs, rhs).1
350-
});
351-
checked
353+
AllOr::Some(valid_bits) => {
354+
checked_constant_array_valid_lanes::<T, Op>(lhs, rhs, valid_bits)
352355
}
353356
}
354357
}
355358

359+
fn checked_array_array_all_lanes<T, Op>(lhs: &[T], rhs: &[T]) -> CheckedValues<T>
360+
where
361+
T: NativePType,
362+
Op: CheckedPrimitiveOp<T>,
363+
{
364+
collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx]))
365+
}
366+
367+
fn checked_array_array_valid_lanes<T, Op>(
368+
lhs: &[T],
369+
rhs: &[T],
370+
valid_bits: &BitBuffer,
371+
) -> CheckedValues<T>
372+
where
373+
T: NativePType,
374+
Op: CheckedPrimitiveOp<T>,
375+
{
376+
let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx]));
377+
378+
checked.failed = checked.failed
379+
&& any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs[idx]).1);
380+
checked
381+
}
382+
383+
fn checked_array_constant_all_lanes<T, Op>(lhs: &[T], rhs: T) -> CheckedValues<T>
384+
where
385+
T: NativePType,
386+
Op: CheckedPrimitiveOp<T>,
387+
{
388+
collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs))
389+
}
390+
391+
fn checked_array_constant_valid_lanes<T, Op>(
392+
lhs: &[T],
393+
rhs: T,
394+
valid_bits: &BitBuffer,
395+
) -> CheckedValues<T>
396+
where
397+
T: NativePType,
398+
Op: CheckedPrimitiveOp<T>,
399+
{
400+
let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs));
401+
402+
checked.failed =
403+
checked.failed && any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs).1);
404+
checked
405+
}
406+
407+
fn checked_constant_array_all_lanes<T, Op>(lhs: T, rhs: &[T]) -> CheckedValues<T>
408+
where
409+
T: NativePType,
410+
Op: CheckedPrimitiveOp<T>,
411+
{
412+
collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx]))
413+
}
414+
415+
fn checked_constant_array_valid_lanes<T, Op>(
416+
lhs: T,
417+
rhs: &[T],
418+
valid_bits: &BitBuffer,
419+
) -> CheckedValues<T>
420+
where
421+
T: NativePType,
422+
Op: CheckedPrimitiveOp<T>,
423+
{
424+
let mut checked = collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx]));
425+
426+
checked.failed =
427+
checked.failed && any_valid_error(rhs.len(), valid_bits, |idx| Op::apply(lhs, rhs[idx]).1);
428+
checked
429+
}
430+
356431
fn collect_all_lanes<T, F>(len: usize, mut value_and_error_at: F) -> CheckedValues<T>
357432
where
358433
T: NativePType,
@@ -375,6 +450,66 @@ where
375450
}
376451
}
377452

453+
fn checked_array_array_one_pass<T, Op>(lhs: &[T], rhs: &[T]) -> CheckedValues<T>
454+
where
455+
T: NativePType,
456+
Op: CheckedPrimitiveOp<T>,
457+
{
458+
checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs[idx]))
459+
}
460+
461+
fn checked_array_array_valid_lanes_one_pass<T, Op>(
462+
lhs: &[T],
463+
rhs: &[T],
464+
valid_bits: &BitBuffer,
465+
) -> CheckedValues<T>
466+
where
467+
T: NativePType,
468+
Op: CheckedPrimitiveOp<T>,
469+
{
470+
checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs[idx]))
471+
}
472+
473+
fn checked_array_constant_one_pass<T, Op>(lhs: &[T], rhs: T) -> CheckedValues<T>
474+
where
475+
T: NativePType,
476+
Op: CheckedPrimitiveOp<T>,
477+
{
478+
checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs))
479+
}
480+
481+
fn checked_array_constant_valid_lanes_one_pass<T, Op>(
482+
lhs: &[T],
483+
rhs: T,
484+
valid_bits: &BitBuffer,
485+
) -> CheckedValues<T>
486+
where
487+
T: NativePType,
488+
Op: CheckedPrimitiveOp<T>,
489+
{
490+
checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs))
491+
}
492+
493+
fn checked_constant_array_one_pass<T, Op>(lhs: T, rhs: &[T]) -> CheckedValues<T>
494+
where
495+
T: NativePType,
496+
Op: CheckedPrimitiveOp<T>,
497+
{
498+
checked_all_lanes(rhs.len(), |idx| Op::checked(lhs, rhs[idx]))
499+
}
500+
501+
fn checked_constant_array_valid_lanes_one_pass<T, Op>(
502+
lhs: T,
503+
rhs: &[T],
504+
valid_bits: &BitBuffer,
505+
) -> CheckedValues<T>
506+
where
507+
T: NativePType,
508+
Op: CheckedPrimitiveOp<T>,
509+
{
510+
checked_valid_lanes(rhs.len(), valid_bits, |idx| Op::checked(lhs, rhs[idx]))
511+
}
512+
378513
// Checked one-pass ops delay early exit until the end of a small block. This
379514
// keeps the loop generic while avoiding a branch-driven exit decision on every
380515
// lane; it is deliberately independent of mask density or input length.

0 commit comments

Comments
 (0)