perf(take): Vectorise bounds check in take_native (-8-10%)#9754
Draft
Dandandan wants to merge 1 commit intoapache:mainfrom
Draft
perf(take): Vectorise bounds check in take_native (-8-10%)#9754Dandandan wants to merge 1 commit intoapache:mainfrom
Dandandan wants to merge 1 commit intoapache:mainfrom
Conversation
The non-null path of `take_native` performed a bounds check per index via `values[index.as_usize()]`. The per-lane branch blocks autovectorisation and dominates the hot loop for primitive take. Reduce each CHUNK=16 indices to their maximum via `fold`+`max` (no short-circuit, so LLVM SIMD-reduces it to two `ldp q` + three `umax.4s` + one `umaxv.4s` on aarch64) and bounds-check the max once per chunk. The panic path is a `#[cold]` helper so `max_idx` does not need to be kept live for format args on the hot path (no stack spill per chunk). Signed index types sign-extend to `usize::MAX` on `as_usize()`, so negative indices still fail the check. Measured on aarch64 (Apple Silicon) with `cargo bench --bench take_kernels`: take i32 512 309 ns → 279 ns (−9.7%) take i32 1024 469 ns → 431 ns (−8.1%) No change to `take` panic semantics (still panics on OOB) or to the null-indices branch. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
The non-null branch of
take_native— the primitivetakegather —bounds-checks each index individually via
values[index.as_usize()].The per-lane branch dominates the loop and blocks autovectorisation.
Reducing a chunk of indices to their maximum with
fold+maxhas noearly exit, so LLVM lowers the whole check to a SIMD max-reduction; on
aarch64 that's two
ldp q+ threeumax.4s+ oneumaxv.4s— asingle bounds check per chunk, then the gather.
max_idxno longer has to stay live on the hot path for the panicformat string because the panic is moved into a
#[cold]helper,which removes a per-chunk stack spill (
str x16, [sp, #16]).Signed index types sign-extend to
usize::MAXonas_usize(), sonegative indices still fail the check; panic behaviour is preserved.
What changes are included in this PR?
One-file change in
arrow-select/src/take.rs:CHUNK = 16chunked max-reduction bounds check#[cold]oobhelper for the panic pathVec::set_len+chunks_exact_mutso thegather is a straight SIMD store (no
push/ capacity-checkoverhead)
Output is written into a
Vec<T>with uninitialised capacity andset_len(len)up front;T: ArrowNativeTypeisCopywith noDrop, and every slot is written beforeoutis read.Are these changes tested?
Yes — covered by existing
take::tests::*, includingtest_take_out_of_bounds_panic.Measured on aarch64 (Apple Silicon) with
cargo bench --bench take_kernels -- "^take i32":Are there any user-facing changes?
No — no API change, panic behaviour on out-of-bounds indices is
preserved.
🤖 Generated with Claude Code