Skip to content

Commit c4adc66

Browse files
committed
more
Signed-off-by: Robert Kruszewski <github@robertk.io>
1 parent c588a9c commit c4adc66

3 files changed

Lines changed: 88 additions & 19 deletions

File tree

‎vortex-array/src/arrays/filter/execute/byte_compress.rs‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -145,11 +145,21 @@ fn filter_chunk_into<T: Copy>(
145145
/// must be valid for writes of every selected element.
146146
pub(super) unsafe fn compress_words<T: Copy>(src: &[T], words: &[u64], dst: *mut T) -> usize {
147147
let mut written = 0;
148-
for (byte_idx, chunk) in src.chunks(8).enumerate() {
149-
let mask_byte = words[byte_idx / 8].to_le_bytes()[byte_idx % 8];
150-
// SAFETY: the mask byte selects only elements of `chunk`, and `dst` has room for every
151-
// selected element.
152-
written += unsafe { compress_byte(chunk, mask_byte, dst.add(written)) };
148+
for (word, src) in words.iter().zip(src.chunks(64)) {
149+
let (chunks, tail) = src.as_chunks::<8>();
150+
let mask_bytes = word.to_le_bytes();
151+
152+
for (chunk, &mask_byte) in chunks.iter().zip(&mask_bytes) {
153+
// SAFETY: the mask byte selects only elements of `chunk`, and `dst` has room for
154+
// every selected element.
155+
written += unsafe { compress_byte(chunk, mask_byte, dst.add(written)) };
156+
}
157+
158+
if !tail.is_empty() {
159+
// SAFETY: the bits past `src.len()` are cleared, so the mask byte selects only
160+
// elements of `tail`.
161+
written += unsafe { compress_byte(tail, mask_bytes[chunks.len()], dst.add(written)) };
162+
}
153163
}
154164
written
155165
}

‎vortex-array/src/arrays/filter/execute/chunked.rs‎

Lines changed: 31 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,9 @@
1414
//! | at most 16 per byte of value width | decode only the selected values |
1515
//! | otherwise | decode to scratch, then compact the scratch |
1616
//!
17-
//! The scratch buffer stays in cache. It is compacted by copying the slices, or with the same
18-
//! SIMD, byte compress and scalar kernels that `filter_buffer` chooses for the chunk's density.
17+
//! The scratch buffer stays in cache. It is compacted by copying the slices, by copying the runs
18+
//! of set bits when they average at least [`MIN_RUN_BYTES`], and otherwise with the same SIMD,
19+
//! byte compress and scalar kernels that `filter_buffer` chooses for the chunk's density.
1920
2021
use std::mem::MaybeUninit;
2122
use std::ptr;
@@ -31,6 +32,7 @@ use super::simd_compress;
3132
use super::simd_compress::SLACK_BYTES;
3233
use super::slice::MaskBits;
3334
use super::slice::compact_by_bitmap;
35+
use super::slice::compact_runs_by_bitmap;
3436
use super::slice::low_bits_mask;
3537
use crate::dtype::NativePType;
3638

@@ -48,6 +50,10 @@ const SPARSE_VALUES_PER_BYTE: usize = 16;
4850
/// The largest number of selected values for which a chunk decodes only those values.
4951
const MAX_SPARSE_VALUES: usize = SPARSE_VALUES_PER_BYTE * 8;
5052

53+
/// Chunks whose runs of selected values average at least this many bytes copy each run. Shorter
54+
/// runs are faster to compact with the SIMD and byte compress kernels.
55+
const MIN_RUN_BYTES: usize = 96;
56+
5157
/// Number of mask bytes read to assemble the words of a chunk at an arbitrary bit offset.
5258
const CHUNK_WORD_BYTES: usize = FILTER_CHUNK_LEN / 8 + 8;
5359

@@ -192,10 +198,9 @@ impl<'a, T: NativePType, D: ChunkDecoder<T>> ChunkedFilter<'a, T, D> {
192198
} else {
193199
let dst = self.dst().cast::<T>();
194200
let src = &self.decode_scratch(chunk_idx)[chunk_offset..][..chunk_len];
195-
let density = selected as f64 / chunk_len as f64;
196201
// SAFETY: the words select only values of `src` and clear any bits past
197202
// `chunk_len`, and the output has room for every selected value plus the slack.
198-
let written = unsafe { compact_words(src, words, density, dst) };
203+
let written = unsafe { compact_words(src, words, selected, dst) };
199204
debug_assert_eq!(written, selected);
200205
self.written += selected;
201206
}
@@ -271,22 +276,33 @@ impl<'a, T: NativePType, D: ChunkDecoder<T>> ChunkedFilter<'a, T, D> {
271276
}
272277
}
273278

274-
/// Copies the values of `src` selected by `words` to `dst` with the kernel that
275-
/// [`filter_buffer`](super::buffer::filter_buffer) chooses for `density`, and returns the number
279+
/// Copies the `selected` values of `src` selected by `words` to `dst` and returns the number
276280
/// copied.
277281
///
282+
/// Runs of set bits that average at least [`MIN_RUN_BYTES`] are copied run by run. Otherwise the kernel is the one that
283+
/// [`filter_buffer`](super::buffer::filter_buffer) chooses for the density.
284+
///
278285
/// # Safety
279286
///
280287
/// `words` must hold `src.len()` bits with any bits past `src.len()` cleared, and `dst` must be
281288
/// valid for writes of every selected value plus [`SLACK_BYTES`].
282-
unsafe fn compact_words<T: Copy>(src: &[T], words: &[u64], density: f64, dst: *mut T) -> usize {
289+
unsafe fn compact_words<T: Copy>(src: &[T], words: &[u64], selected: usize, dst: *mut T) -> usize {
283290
let bits = MaskBits::Words {
284291
words,
285292
len: src.len(),
286293
};
294+
// A run that crosses a word boundary counts once in each word, as the run walk copies it.
295+
let runs = words
296+
.iter()
297+
.map(|&word| (word & !(word << 1)).count_ones() as usize)
298+
.sum::<usize>();
299+
let density = selected as f64 / src.len() as f64;
300+
287301
// SAFETY: forwarded from the caller contract.
288302
unsafe {
289-
if let Some(written) = simd_compress::compress_bits(src, bits, density, dst) {
303+
if selected * size_of::<T>() >= MIN_RUN_BYTES * runs {
304+
compact_runs_by_bitmap(src, bits, dst)
305+
} else if let Some(written) = simd_compress::compress_bits(src, bits, density, dst) {
290306
written
291307
} else if density >= byte_compress_density_threshold::<T>() {
292308
byte_compress::compress_words(src, words, dst)
@@ -422,19 +438,20 @@ mod tests {
422438

423439
/// Runs long enough that the cached slices are used rather than the bitmap.
424440
#[rstest]
425-
fn filter_chunked_by_slices_matches_expected(
441+
fn filter_chunked_runs_matches_expected(
426442
#[values(9, 12, 200, 3000)] run_len: usize,
427-
#[values(16, 64, 1500)] gap: usize,
443+
#[values(1, 16, 64, 1500)] gap: usize,
428444
#[values(0, 3, 1000)] offset: usize,
445+
#[values(false, true)] slices: bool,
429446
) {
430447
let bits = (0..5000)
431448
.map(|i| i % (run_len + gap) < run_len)
432449
.collect::<BitBuffer>();
433450

434-
check::<u8>(&bits, offset, true);
435-
check::<u16>(&bits, offset, true);
436-
check::<u32>(&bits, offset, true);
437-
check::<u64>(&bits, offset, true);
451+
check::<u8>(&bits, offset, slices);
452+
check::<u16>(&bits, offset, slices);
453+
check::<u32>(&bits, offset, slices);
454+
check::<u64>(&bits, offset, slices);
438455
}
439456

440457
#[rstest]

‎vortex-array/src/arrays/filter/execute/slice.rs‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,48 @@ pub(super) unsafe fn compact_by_bitmap<T: Copy>(
145145
write_pos
146146
}
147147

148+
/// Copy the elements of `src` selected by `bits` to `dst` one run of set bits at a time, and
149+
/// return the number copied.
150+
///
151+
/// This is faster than [`compact_by_bitmap`] when the selected elements form long runs.
152+
///
153+
/// # Safety
154+
///
155+
/// `bits` must select only elements of `src`, and `dst` must be valid for writes of every
156+
/// selected element.
157+
// Inlining this walk next to `compact_by_bitmap` makes the bit walk 1.5x slower on Apple M4.
158+
#[inline(never)]
159+
pub(super) unsafe fn compact_runs_by_bitmap<T: Copy>(
160+
src: &[T],
161+
bits: MaskBits<'_>,
162+
dst: *mut T,
163+
) -> usize {
164+
let src_ptr = src.as_ptr();
165+
let mut write_pos = 0;
166+
167+
for_each_mask_word(bits, |mut word, word_start, _| {
168+
while word != 0 {
169+
let run_start = word.trailing_zeros() as usize;
170+
let run_len = (!(word >> run_start)).trailing_zeros() as usize;
171+
// SAFETY: the run lies within the selected elements of `src`, and `dst` has room for
172+
// every selected element.
173+
unsafe {
174+
ptr::copy_nonoverlapping(
175+
src_ptr.add(word_start + run_start),
176+
dst.add(write_pos),
177+
run_len,
178+
)
179+
};
180+
write_pos += run_len;
181+
182+
// Adding the lowest bit of the run carries through the run and clears it.
183+
word &= word.wrapping_add(1 << run_start);
184+
}
185+
});
186+
187+
write_pos
188+
}
189+
148190
/// Filter a slice by a set of strictly increasing indices.
149191
pub(super) fn filter_slice_by_indices<T: Copy>(
150192
slice: &[T],

0 commit comments

Comments
 (0)