Skip to content

Commit b7fb836

Browse files
connortsui20connortsui20
authored
Perf: optimize take_scalar (#5723)
This change optimizes the scalar `take` compute function over slices. It is pretty crazy how the compiler was not able to figure this out on its own... We can look at the codspeed numbers but on my machine this is ~35% faster: Original: <details> ``` ├─ pvector_take_uniform │ │ │ │ │ │ ├─ 16 │ │ │ │ │ │ │ ├─ 1000 404.8 ns │ 6.394 µs │ 409.8 ns │ 413.1 ns │ 100000 │ 400000 │ │ ├─ 10000 3.689 µs │ 34.39 µs │ 3.709 µs │ 3.746 µs │ 100000 │ 100000 │ │ ╰─ 100000 36.43 µs │ 124 µs │ 36.63 µs │ 37 µs │ 100000 │ 100000 │ ├─ 256 │ │ │ │ │ │ │ ├─ 1000 404.8 ns │ 20.97 µs │ 409.8 ns │ 419.7 ns │ 100000 │ 400000 │ │ ├─ 10000 3.679 µs │ 81.82 µs │ 3.699 µs │ 3.78 µs │ 100000 │ 100000 │ │ ╰─ 100000 36.45 µs │ 122.8 µs │ 36.54 µs │ 36.95 µs │ 100000 │ 100000 │ ├─ 2048 │ │ │ │ │ │ │ ├─ 1000 409.8 ns │ 30.72 µs │ 429.8 ns │ 433.4 ns │ 100000 │ 100000 │ │ ├─ 10000 3.689 µs │ 21.1 µs │ 3.699 µs │ 3.718 µs │ 100000 │ 100000 │ │ ╰─ 100000 36.51 µs │ 124.3 µs │ 36.79 µs │ 37.2 µs │ 100000 │ 100000 │ ╰─ 8192 │ │ │ │ │ │ ├─ 1000 419.8 ns │ 5.489 µs │ 434.8 ns │ 437.4 ns │ 100000 │ 200000 │ ├─ 10000 3.699 µs │ 38.87 µs │ 3.749 µs │ 3.793 µs │ 100000 │ 100000 │ ╰─ 100000 36.55 µs │ 122.1 µs │ 36.95 µs │ 37.19 µs │ 100000 │ 100000 ╰─ pvector_take_zipfian │ │ │ │ │ ├─ 16 │ │ │ │ │ │ ├─ 1000 404.8 ns │ 2.012 µs │ 407.3 ns │ 409.6 ns │ 100000 │ 400000 │ ├─ 10000 3.689 µs │ 30.16 µs │ 3.719 µs │ 3.729 µs │ 100000 │ 100000 │ ╰─ 100000 36.53 µs │ 125.6 µs │ 36.78 µs │ 37.06 µs │ 100000 │ 100000 ├─ 256 │ │ │ │ │ │ ├─ 1000 409.8 ns │ 4.949 µs │ 414.8 ns │ 415.4 ns │ 100000 │ 200000 │ ├─ 10000 3.689 µs │ 29.16 µs │ 3.719 µs │ 3.731 µs │ 100000 │ 100000 │ ╰─ 100000 36.52 µs │ 122.6 µs │ 36.77 µs │ 37 µs │ 100000 │ 100000 ├─ 2048 │ │ │ │ │ │ ├─ 1000 399.8 ns │ 6.302 µs │ 404.8 ns │ 428.2 ns │ 100000 │ 400000 │ ├─ 10000 3.689 µs │ 25.36 µs │ 3.699 µs │ 3.736 µs │ 100000 │ 100000 │ ╰─ 100000 36.5 µs │ 121.6 µs │ 36.72 µs │ 36.96 µs │ 100000 │ 100000 ╰─ 8192 │ │ │ │ │ ├─ 1000 407.3 ns │ 1.914 µs │ 412.3 ns │ 415.2 ns │ 100000 │ 400000 ├─ 10000 3.689 µs │ 13.83 µs │ 3.729 µs │ 3.744 µs │ 100000 │ 100000 ╰─ 100000 36.47 µs │ 126.8 µs │ 37.41 µs │ 37.79 µs │ 100000 │ 100000 ``` </details> New: <details> ``` take_primitive fastest │ slowest │ median │ mean │ samples │ iters ├─ pvector_take_uniform │ │ │ │ │ │ ├─ 16 │ │ │ │ │ │ │ ├─ 1000 279.8 ns │ 15.9 µs │ 299.8 ns │ 300.4 ns │ 100000 │ 100000 │ │ ├─ 10000 2.529 µs │ 22.95 µs │ 2.569 µs │ 2.592 µs │ 100000 │ 100000 │ │ ╰─ 100000 23.32 µs │ 115.6 µs │ 23.66 µs │ 23.95 µs │ 100000 │ 100000 │ ├─ 256 │ │ │ │ │ │ │ ├─ 1000 257.3 ns │ 1.302 µs │ 264.8 ns │ 267.2 ns │ 100000 │ 400000 │ │ ├─ 10000 2.519 µs │ 42.61 µs │ 2.569 µs │ 2.606 µs │ 100000 │ 100000 │ │ ╰─ 100000 23.59 µs │ 117.2 µs │ 23.71 µs │ 24.03 µs │ 100000 │ 100000 │ ├─ 2048 │ │ │ │ │ │ │ ├─ 1000 262.3 ns │ 3.857 µs │ 267.3 ns │ 270.4 ns │ 100000 │ 400000 │ │ ├─ 10000 2.559 µs │ 19.81 µs │ 2.599 µs │ 2.631 µs │ 100000 │ 100000 │ │ ╰─ 100000 23.48 µs │ 111 µs │ 23.62 µs │ 23.99 µs │ 100000 │ 100000 │ ╰─ 8192 │ │ │ │ │ │ ├─ 1000 292.3 ns │ 6.382 µs │ 302.3 ns │ 304.2 ns │ 100000 │ 400000 │ ├─ 10000 2.749 µs │ 21.31 µs │ 2.799 µs │ 2.819 µs │ 100000 │ 100000 │ ╰─ 100000 25.28 µs │ 118.3 µs │ 25.56 µs │ 25.9 µs │ 100000 │ 100000 ╰─ pvector_take_zipfian │ │ │ │ │ ├─ 16 │ │ │ │ │ │ ├─ 1000 257.3 ns │ 5.714 µs │ 264.8 ns │ 267.7 ns │ 100000 │ 400000 │ ├─ 10000 2.519 µs │ 20.1 µs │ 2.569 µs │ 2.593 µs │ 100000 │ 100000 │ ╰─ 100000 23.43 µs │ 105.6 µs │ 23.6 µs │ 23.87 µs │ 100000 │ 100000 ├─ 256 │ │ │ │ │ │ ├─ 1000 259.8 ns │ 1.259 µs │ 264.8 ns │ 267.7 ns │ 100000 │ 400000 │ ├─ 10000 2.509 µs │ 36.64 µs │ 2.569 µs │ 2.763 µs │ 100000 │ 100000 │ ╰─ 100000 23.2 µs │ 107 µs │ 23.59 µs │ 23.87 µs │ 100000 │ 100000 ├─ 2048 │ │ │ │ │ │ ├─ 1000 259.8 ns │ 1.957 µs │ 267.3 ns │ 275.8 ns │ 100000 │ 400000 │ ├─ 10000 2.569 µs │ 9.469 µs │ 2.609 µs │ 2.644 µs │ 100000 │ 100000 │ ╰─ 100000 23.72 µs │ 109.6 µs │ 23.99 µs │ 24.27 µs │ 100000 │ 100000 ╰─ 8192 │ │ │ │ │ ├─ 1000 269.8 ns │ 26.15 µs │ 284.8 ns │ 289.4 ns │ 100000 │ 200000 ├─ 10000 2.709 µs │ 8.879 µs │ 2.739 µs │ 2.751 µs │ 100000 │ 100000 ╰─ 100000 24.27 µs │ 107.2 µs │ 24.53 µs │ 24.74 µs │ 100000 │ 100000 ``` </details> Signed-off-by: Connor Tsui <[email protected]>
1 parent afe1196 commit b7fb836

File tree

1 file changed

+17
-2
lines changed
  • vortex-compute/src/take/slice

1 file changed

+17
-2
lines changed

‎vortex-compute/src/take/slice/mod.rs‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
//! Take function implementations on slices.
55
66
use vortex_buffer::Buffer;
7+
use vortex_buffer::BufferMut;
78
use vortex_dtype::UnsignedPType;
89

910
use crate::take::Take;
@@ -43,7 +44,21 @@ impl<T: Copy, I: UnsignedPType> Take<[I]> for &[T] {
4344
unused,
4445
reason = "Compiler may see this as unused based on enabled features"
4546
)]
46-
#[inline]
4747
fn take_scalar<T: Copy, I: UnsignedPType>(buffer: &[T], indices: &[I]) -> Buffer<T> {
48-
indices.iter().map(|idx| buffer[idx.as_()]).collect()
48+
// NB: The simpler `indices.iter().map(|idx| buff1er[idx.as_()]).collect()` generates suboptimal
49+
// assembly where the buffer length is repeatedly loaded from the stack on each iteration.
50+
51+
let mut result = BufferMut::with_capacity(indices.len());
52+
let ptr = result.spare_capacity_mut().as_mut_ptr().cast::<T>();
53+
54+
// This explicit loop with pointer writes keeps the length in a register and avoids per-element
55+
// capacity checks from `push()`.
56+
for (i, idx) in indices.iter().enumerate() {
57+
// SAFETY: We reserved `indices.len()` capacity, so `ptr.add(i)` is valid.
58+
unsafe { ptr.add(i).write(buffer[idx.as_()]) };
59+
}
60+
61+
// SAFETY: We just wrote exactly `indices.len()` elements.
62+
unsafe { result.set_len(indices.len()) };
63+
result.freeze()
4964
}

0 commit comments

Comments
 (0)