Skip to content

Commit 6c2eb22

Browse files
authored
perf(array): pack Boolean output during RowFn dense retry (#9986)
Stacked on #9988, which adds the CI benchmark without changing production code. ## Summary Routes deferred Boolean dense attempts through the packed collector, preserving the requested multiversioning flag. Keeps the source and failure accumulator in one borrowed state, and inlines the collector's scalar tail so LLVM can retain the alias information needed to vectorize the row loop. Decode errors remain terminal, and null-only failures still retry valid rows. The earlier x86 regressions came from scalarized collector loops. Optimized Linux IR and the native CI executables confirm that these loops now vectorize. ## Changes Native CI medians in microseconds per eight 16,384-row batches, with partial validity and the multiversioned collector: | Target | Columns before → after | Constant LHS before → after | | --- | ---: | ---: | | AVX2 | 111.0 → 44.53 | 217.1 → 34.58 | | AVX-512 | 42.81 → 45.07 | 206.9 → 37.25 | | NEON | 68.32 → 71.01 | 173.3 → 57.62 | Constant RHS also improves on every target. Small column-only movements remain uncertain after one paired run. All final cases are below 1 ms, with a maximum printed sample of 555.8 µs. Passed 175 RowFn tests, 14 Boolean packing tests, and scoped Clippy. The CodSpeed check flags no Boolean retry regressions, but remains red on three other, unattributed benchmark alerts. This remains a draft. [Four-stage comparison, compiler evidence, and raw results](https://github.com/vortex-data/vortex/blob/ct/row-fn-performance-evidence/research/row-fn-engine/performance/ci-bool-capture.md). Evidence stays on the separate branch. Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 2881302 commit 6c2eb22

5 files changed

Lines changed: 262 additions & 5 deletions

File tree

‎vortex-array/src/scalar_fn/unstable/row/batch/tests.rs‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@ use crate::arrays::ConstantArray;
2525
use crate::arrays::ExtensionArray;
2626
use crate::arrays::FixedSizeListArray;
2727
use crate::arrays::PrimitiveArray;
28+
use crate::arrays::VarBinViewArray;
29+
use crate::arrays::varbinview::BinaryView;
2830
use crate::assert_arrays_eq;
31+
use crate::buffer::BufferHandle;
2932
use crate::dtype::DType;
3033
use crate::dtype::NativePType;
3134
use crate::dtype::Nullability;
@@ -43,6 +46,8 @@ use crate::scalar_fn::unstable::row::OutputElement;
4346
use crate::scalar_fn::unstable::row::OutputSink;
4447
use crate::scalar_fn::unstable::row::RowFn;
4548
use crate::scalar_fn::unstable::row::RowVisitor;
49+
use crate::scalar_fn::unstable::row::Utf8Column;
50+
use crate::scalar_fn::unstable::row::execute::execute_bool_dense_attempt;
4651
use crate::scalar_fn::unstable::row::execute_rows;
4752
use crate::scalar_fn::unstable::row::row_fn_return_dtype;
4853
use crate::validity::Validity;
@@ -83,6 +88,20 @@ struct PackedGreaterThan<const MULTIVERSIONED: bool>;
8388
#[derive(Clone)]
8489
struct DeferredGreaterThan<const MULTIVERSIONED: bool>;
8590

91+
#[derive(Clone, Copy)]
92+
enum BooleanInputShape {
93+
Columns,
94+
ConstantLhs,
95+
ConstantRhs,
96+
}
97+
98+
#[derive(Clone, Copy)]
99+
enum BooleanValidity {
100+
AllValid,
101+
AllNull,
102+
Partial,
103+
}
104+
86105
#[derive(Clone)]
87106
struct ValidOnlyPositive;
88107

@@ -866,6 +885,100 @@ fn test_deferred_bool_output_handles_partial_constants() -> VortexResult<()> {
866885
Ok(())
867886
}
868887

888+
#[rstest]
889+
#[case::empty(0)]
890+
#[case::one(1)]
891+
#[case::word_tail(63)]
892+
#[case::word(64)]
893+
#[case::word_remainder(65)]
894+
#[case::multiword(130)]
895+
fn test_deferred_bool_sliced_validity_and_constants(
896+
#[case] len: usize,
897+
#[values(false, true)] multiversioned: bool,
898+
#[values(
899+
BooleanInputShape::Columns,
900+
BooleanInputShape::ConstantLhs,
901+
BooleanInputShape::ConstantRhs
902+
)]
903+
input_shape: BooleanInputShape,
904+
#[values(
905+
BooleanValidity::AllValid,
906+
BooleanValidity::AllNull,
907+
BooleanValidity::Partial
908+
)]
909+
validity_pattern: BooleanValidity,
910+
) -> VortexResult<()> {
911+
let valid = |index: usize| match validity_pattern {
912+
BooleanValidity::AllValid => true,
913+
BooleanValidity::AllNull => false,
914+
BooleanValidity::Partial => !index.is_multiple_of(3),
915+
};
916+
let validity = Validity::Array(BoolArray::from_iter((0..len + 2).map(valid)).into_array());
917+
let varying = PrimitiveArray::new(
918+
(0..len + 2)
919+
.map(|index| if valid(index) { index as i64 } else { i64::MIN })
920+
.collect::<Vec<_>>(),
921+
validity,
922+
)
923+
.into_array()
924+
.slice(1..len + 1)?;
925+
let constant = ConstantArray::new(32_i64, len).into_array();
926+
let other = PrimitiveArray::from_iter(std::iter::repeat_n(32_i64, len)).into_array();
927+
let (lhs, rhs) = match input_shape {
928+
BooleanInputShape::Columns => (varying, other),
929+
BooleanInputShape::ConstantLhs => (constant, varying),
930+
BooleanInputShape::ConstantRhs => (varying, constant),
931+
};
932+
let args = VecExecutionArgs::new(vec![lhs, rhs], len);
933+
let mut ctx = array_session().create_execution_ctx();
934+
let actual = if multiversioned {
935+
execute_rows(&DeferredGreaterThan::<true>, &EmptyOptions, &args, &mut ctx)?
936+
} else {
937+
execute_rows(
938+
&DeferredGreaterThan::<false>,
939+
&EmptyOptions,
940+
&args,
941+
&mut ctx,
942+
)?
943+
};
944+
let expected = BoolArray::from_iter((1..len + 1).map(|index| {
945+
valid(index).then_some(if matches!(input_shape, BooleanInputShape::ConstantLhs) {
946+
32 > index as i64
947+
} else {
948+
index as i64 > 32
949+
})
950+
}))
951+
.into_array();
952+
assert_arrays_eq!(&actual, &expected, &mut ctx);
953+
Ok(())
954+
}
955+
956+
#[test]
957+
fn test_deferred_bool_decode_error_is_terminal() {
958+
let views = Buffer::from(vec![BinaryView::make_view(&[0xff], 0, 0)]);
959+
let input = VarBinViewArray::new_handle(
960+
BufferHandle::new_host(views.into_byte_buffer()),
961+
Default::default(),
962+
DType::Utf8(Nullability::NonNullable),
963+
Validity::NonNullable,
964+
)
965+
.into_array();
966+
let args = VecExecutionArgs::new(vec![input], 1);
967+
let mut ctx = array_session().create_execution_ctx();
968+
969+
let result = execute_bool_dense_attempt::<(Utf8Column,), (), bool, true>(
970+
&args,
971+
&mut ctx,
972+
|_| (),
973+
|&(), (value,)| (!value.is_empty(), false),
974+
|_| Ok(()),
975+
);
976+
assert!(
977+
result.is_err(),
978+
"a decode error must not become a deferred row error"
979+
);
980+
}
981+
869982
#[test]
870983
fn test_deferred_bool_output_reports_valid_row_failure() -> VortexResult<()> {
871984
let lhs = PrimitiveArray::from_iter([1_i64, i64::MIN, -1]).into_array();

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pub(super) use owned::execute_owned_infallible_valid_rows;
1515
pub(super) use owned::execute_owned_valid_rows;
1616

1717
mod packed_bool;
18+
pub(super) use packed_bool::execute_bool_dense_attempt;
1819
pub(super) use packed_bool::execute_owned_bool;
1920
pub(super) use packed_bool::execute_owned_infallible_bool;
2021

‎vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs‎

Lines changed: 94 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@
55
//!
66
//! These executors share the indexed input decoding used by owned outputs, but construct a
77
//! canonical [`BoolArray`] without first collecting one byte per row.
8+
//!
9+
//! [`execute_owned_bool`] reports rejected failure evidence as a terminal error.
10+
//! [`execute_bool_dense_attempt`] hands it back as a [`DenseAttempt::DeferredError`] so that batch
11+
//! execution can retry over valid rows. The two keep separate copies of the row loop on purpose.
812
913
use std::mem::size_of;
1014

@@ -13,6 +17,7 @@ use vortex_compute::lane_kernels::IndexedSource;
1317
use vortex_error::VortexResult;
1418
use vortex_error::vortex_bail;
1519

20+
use super::DenseAttempt;
1621
use crate::ArrayRef;
1722
use crate::ExecutionCtx;
1823
use crate::IntoArray;
@@ -82,14 +87,15 @@ where
8287

8388
let mut failure = Fail::default();
8489
let collect = |index| {
85-
// SAFETY: Both collectors only invoke this closure with `index < row_count`, and the decoded
86-
// source was constructed with exactly `row_count` rows.
90+
// SAFETY: Both collectors only invoke this closure with `index < row_count`, and the
91+
// decoded source was constructed with exactly `row_count` rows.
8792
let elements = unsafe { source.get_unchecked(index) };
8893
let (value, row_failure) = apply(&prepared, elements);
8994
failure |= row_failure;
9095

9196
value
9297
};
98+
9399
let values = if MULTIVERSIONED {
94100
BitBuffer::collect_bool_multiversioned(row_count, collect)
95101
} else {
@@ -100,3 +106,89 @@ where
100106

101107
Ok(BoolArray::new(values, Validity::NonNullable).into_array())
102108
}
109+
110+
/// Pack a dense Boolean attempt and report rejected failure evidence to the batch executor.
111+
///
112+
/// Returns the same terminal-versus-deferred split as
113+
/// [`execute_owned_dense_attempt`](super::execute_owned_dense_attempt), which owns that contract.
114+
pub(crate) fn execute_bool_dense_attempt<Args, Prepared, Fail, const MULTIVERSIONED: bool>(
115+
args: &dyn ExecutionArgs,
116+
ctx: &mut ExecutionCtx,
117+
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
118+
apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail),
119+
finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
120+
) -> VortexResult<DenseAttempt>
121+
where
122+
Args: IndexedElementTuple,
123+
Fail: FailureEvidence,
124+
{
125+
const {
126+
assert!(
127+
size_of::<Fail>() <= size_of::<bool>(),
128+
"failure evidence must be no wider than the value, or it bounds the vector width"
129+
)
130+
};
131+
132+
// Keep this dense row loop separate from `execute_owned_bool`. Factoring their shared state
133+
// into a helper changes LLVM's optimized dense kernel even when the helper is inlined.
134+
let columns = Args::decode(args, ctx)?;
135+
let prepared = prepare(Args::const_values(&columns));
136+
let row_count = args.row_count();
137+
138+
let Some(source) = decoded_source::<Args>(&columns, row_count) else {
139+
vortex_bail!("a decoded row input does not address exactly {row_count} rows");
140+
};
141+
142+
// NB: The collector must capture one borrow of the whole state. That shape is load-bearing for
143+
// code generation, not a borrow checker workaround. See the `DeferredBoolState` documentation.
144+
let mut state = DeferredBoolState {
145+
source,
146+
prepared,
147+
apply,
148+
failure: Fail::default(),
149+
};
150+
let state_ref = &mut state;
151+
152+
let collect = move |index| {
153+
let state = &mut *state_ref;
154+
155+
// SAFETY: Both collectors only invoke this closure with `index < row_count`, and the
156+
// decoded source was constructed with exactly `row_count` rows.
157+
let elements = unsafe { state.source.get_unchecked(index) };
158+
let (value, row_failure) = (state.apply)(&state.prepared, elements);
159+
state.failure |= row_failure;
160+
161+
value
162+
};
163+
164+
let values = if MULTIVERSIONED {
165+
BitBuffer::collect_bool_multiversioned(row_count, collect)
166+
} else {
167+
BitBuffer::collect_bool(row_count, collect)
168+
};
169+
170+
match finish_failure(state.failure) {
171+
Ok(()) => Ok(DenseAttempt::Values(
172+
BoolArray::new(values, Validity::NonNullable).into_array(),
173+
)),
174+
Err(error) => Ok(DenseAttempt::DeferredError(error)),
175+
}
176+
}
177+
178+
/// Row-loop state for a deferred Boolean computation: the decoded input, the prepared batch
179+
/// constants, the per-row callback, and the accumulated failure evidence.
180+
///
181+
/// The collector closure captures this struct through one mutable borrow. It does not capture the
182+
/// four parts separately. The single borrow gives the out-of-line collector one exclusive pointer
183+
/// to its state, and the vectorizer needs that to keep the multiversioned word loop packed.
184+
/// Separate captures lose the same aliasing information, and the loop stays scalar.
185+
///
186+
/// The capture shape is therefore load-bearing for code generation. The measurements come from
187+
/// the `row_fn_bool_retry` benchmark, which builds in the `bench` profile with 16 codegen units
188+
/// and no LTO. Recheck that benchmark and its optimized IR before you change the captures.
189+
struct DeferredBoolState<Source, Prepared, Apply, Fail> {
190+
source: Source,
191+
prepared: Prepared,
192+
apply: Apply,
193+
failure: Fail,
194+
}

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ use crate::scalar_fn::unstable::row::RowFn;
3131
use crate::scalar_fn::unstable::row::SinkResult;
3232
use crate::scalar_fn::unstable::row::batch::BorrowedRowFnArgs;
3333
use crate::scalar_fn::unstable::row::execute::DenseAttempt;
34+
use crate::scalar_fn::unstable::row::execute::execute_bool_dense_attempt;
3435
use crate::scalar_fn::unstable::row::execute::execute_owned_dense_attempt;
3536
use crate::scalar_fn::unstable::row::execute::execute_owned_infallible;
3637
use crate::scalar_fn::unstable::row::execute::execute_sink;
@@ -152,4 +153,31 @@ impl<F: RowFn> RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> {
152153
finish_failure,
153154
)
154155
}
156+
157+
fn visit_prepared_deferred_bool<Args, Prepared, Fail, const MULTIVERSIONED: bool>(
158+
self,
159+
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
160+
apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail),
161+
finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
162+
) -> VortexResult<Self::VisitResult>
163+
where
164+
Args: IndexedElementTuple,
165+
Fail: FailureEvidence,
166+
{
167+
const { assert_deferred_visit_contract::<F, Args, bool, Fail>() };
168+
let visited = BatchPlan::new(
169+
validate_owned_visit::<Args, bool>(self.args.dtypes())?,
170+
self.output_dtype,
171+
RowPolicy::for_deferred_output::<Args>(),
172+
)?;
173+
self.args.plan().ensure_reproduced_by(&visited)?;
174+
175+
execute_bool_dense_attempt::<Args, Prepared, Fail, MULTIVERSIONED>(
176+
self.args,
177+
self.ctx,
178+
prepare,
179+
apply,
180+
finish_failure,
181+
)
182+
}
155183
}

‎vortex-buffer/src/bit/pack.rs‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,9 @@
2121
//! simple enough that the per-level duplication and its `#[target_feature]` call boundary pay
2222
//! off.
2323
//!
24-
//! The bit-at-a-time loop lives on as [`collect_bool_word_scalar`], used for tail chunks and as
25-
//! the reference implementation for tests and benchmarks.
24+
//! The bit-at-a-time loop lives on as [`collect_bool_word_scalar`], the reference implementation
25+
//! for tests and benchmarks. The word loop packs its tail chunk with [`collect_bool_word_tail`],
26+
//! its own copy of that loop.
2627
2728
/// Packs up to 64 boolean values into a little-endian `u64` word one bit at a time.
2829
///
@@ -172,10 +173,32 @@ where
172173

173174
if remainder != 0 {
174175
let offset = full * 64;
175-
words[full] = collect_bool_word_scalar(remainder, |bit_idx| f(offset + bit_idx));
176+
words[full] = collect_bool_word_tail(remainder, |bit_idx| f(offset + bit_idx));
176177
}
177178
}
178179

180+
/// Tail chunk of [`collect_bool_words_with`]: a copy of [`collect_bool_word_scalar`] that always
181+
/// inlines. `len` **must** be at most 64.
182+
///
183+
/// The word loop must keep the borrowed callback inline. An out-of-line tail lets the callback
184+
/// state escape, and the full words packed by the same loop then fail to vectorize.
185+
/// [`collect_bool_word_scalar`] stays an ordinary `#[inline]` function so that this constraint does
186+
/// not reach its public callers, which include the scalar benchmark baselines.
187+
#[expect(clippy::inline_always)]
188+
#[inline(always)]
189+
fn collect_bool_word_tail<F>(len: usize, mut f: F) -> u64
190+
where
191+
F: FnMut(usize) -> bool,
192+
{
193+
debug_assert!(len <= 64, "cannot pack {len} bits into a u64 word");
194+
195+
let mut packed = 0;
196+
for bit_idx in 0..len {
197+
packed |= (f(bit_idx) as u64) << bit_idx;
198+
}
199+
packed
200+
}
201+
179202
/// SSE2 copy of the [`collect_bool_words`](crate::bit::collect_bool_words) word loop.
180203
///
181204
/// `words` must hold at least `len.div_ceil(64)` entries and `f` is invoked with `0..len`.

0 commit comments

Comments
 (0)