Skip to content

Commit 96cfdf2

Browse files
authored
chore: simplify dot implementation to use auto-vectorization (#2645)
This change makes the auto-vectorization version of dot(f32) as fast as manually written SIMD. Run benchmarks via ``` export RUSTFLAGS="-C target-cpu=native" git checkout main cargo bench --bench dot -- --save-baseline dot_main f32 git checkout lei/simplify_dot cargo bench --bench dot -- --baseline dot_main f32 ``` On Macbook M2 Max ``` Dot(f32, auto-vectorization) time: [88.812 ms 89.654 ms 90.306 ms] change: [-2.5819% -1.6876% -0.6964%] (p = 0.01 < 0.10) Change within noise threshold. ``` AMD 5900X ``` Dot(f32, auto-vectorization) time: [172.50 ms 176.41 ms 179.41 ms] change: [-2.3545% +0.6133% +3.5448%] (p = 0.69 > 0.10) No change in performance detected. ``` Intel Sapphire ``` Dot(f32, auto-vectorization) time: [331.36 ms 331.62 ms 331.93 ms] change: [-2.3160% -1.1226% -0.3451%] (p = 0.04 < 0.10) Change within noise threshold. ``` Graviton3 ``` Benchmarking Dot(f32, auto-vectorization): Warming up for 3.0000 s Warning: Unable to complete 10 samples in 5.0s. You may wish to increase target time to 8.8s or enable flat sampling. Dot(f32, auto-vectorization) time: [160.62 ms 160.70 ms 160.76 ms] change: [-1.1157% -0.6868% -0.2951%] (p = 0.00 < 0.10) Change within noise threshold. Found 1 outliers among 10 measurements (10.00%) 1 (10.00%) low mild ```
1 parent c7dce32 commit 96cfdf2

1 file changed

Lines changed: 1 addition & 54 deletions

File tree

  • rust/lance-linalg/src/distance

rust/lance-linalg/src/distance/dot.rs

Lines changed: 1 addition & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,6 @@ use lance_core::utils::cpu::SimdSupport;
1919
use lance_core::utils::cpu::SIMD_SUPPORT;
2020
use num_traits::{real::Real, AsPrimitive, Num};
2121

22-
use crate::simd::{
23-
f32::{f32x16, f32x8},
24-
SIMD,
25-
};
2622
use crate::Result;
2723

2824
/// Default implementation of dot product.
@@ -142,56 +138,7 @@ impl Dot for f16 {
142138
impl Dot for f32 {
143139
#[inline]
144140
fn dot(x: &[Self], y: &[Self]) -> f32 {
145-
// Manually unrolled 8 times to get enough registers.
146-
// TODO: avx512 can unroll more
147-
let x_unrolled_chunks = x.chunks_exact(64);
148-
let y_unrolled_chunks = y.chunks_exact(64);
149-
150-
// 8 float32 SIMD
151-
let x_aligned_chunks = x_unrolled_chunks.remainder().chunks_exact(8);
152-
let y_aligned_chunks = y_unrolled_chunks.remainder().chunks_exact(8);
153-
154-
let sum = if x_aligned_chunks.remainder().is_empty() {
155-
0.0
156-
} else {
157-
debug_assert_eq!(
158-
x_aligned_chunks.remainder().len(),
159-
y_aligned_chunks.remainder().len()
160-
);
161-
x_aligned_chunks
162-
.remainder()
163-
.iter()
164-
.zip(y_aligned_chunks.remainder().iter())
165-
.map(|(&x, &y)| x * y)
166-
.sum()
167-
};
168-
169-
let mut sum8 = f32x8::zeros();
170-
x_aligned_chunks
171-
.zip(y_aligned_chunks)
172-
.for_each(|(x_chunk, y_chunk)| unsafe {
173-
let x1 = f32x8::load_unaligned(x_chunk.as_ptr());
174-
let y1 = f32x8::load_unaligned(y_chunk.as_ptr());
175-
sum8 += x1 * y1;
176-
});
177-
178-
let mut sum16 = f32x16::zeros();
179-
x_unrolled_chunks
180-
.zip(y_unrolled_chunks)
181-
.for_each(|(x, y)| unsafe {
182-
let x1 = f32x16::load_unaligned(x.as_ptr());
183-
let x2 = f32x16::load_unaligned(x.as_ptr().add(16));
184-
let x3 = f32x16::load_unaligned(x.as_ptr().add(32));
185-
let x4 = f32x16::load_unaligned(x.as_ptr().add(48));
186-
187-
let y1 = f32x16::load_unaligned(y.as_ptr());
188-
let y2 = f32x16::load_unaligned(y.as_ptr().add(16));
189-
let y3 = f32x16::load_unaligned(y.as_ptr().add(32));
190-
let y4 = f32x16::load_unaligned(y.as_ptr().add(48));
191-
192-
sum16 += (x1 * y1 + x2 * y2) + (x3 * y3 + x4 * y4);
193-
});
194-
sum16.reduce_sum() + sum8.reduce_sum() + sum
141+
dot_scalar::<Self, Self, 16>(x, y)
195142
}
196143
}
197144

0 commit comments

Comments
 (0)