Skip to content

Commit 18ac42b

Browse files
authored
Speed up AVX-512 byte and fixed-width array scans (#55)
* Add repeatable x86 SIMD benchmark harness Capture CPU metadata and run Criterion plus end-to-end Python benchmarks three times on the same AVX2 or AVX-512 machine. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 * Expand x86 threshold benchmark coverage Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 * Benchmark native x86 byte dispatch Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 * Try masked AVX-512 for short byte inputs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 * Try AVX-512 cross-record array scans Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 * Preserve first-record scanner short circuit Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 * Fix AVX scanner exact-match test threshold Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 * Document measured AVX-512 speedups Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007 --------- Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007
1 parent fa1a087 commit 18ac42b

6 files changed

Lines changed: 963 additions & 11 deletions

File tree

README.rst

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ Lastly, I wanted to minimize dependencies, meaning you do not need to install
6565
As of v3.0.0, ``hexhamming`` is written in Rust using `PyO3 <https://pyo3.rs>`_
6666
and `maturin <https://www.maturin.rs>`_, providing memory safety, GIL release
6767
during computation, and free-threaded Python support while maintaining the same
68-
SIMD-accelerated performance (SSE4.1, AVX2, NEON).
68+
SIMD-accelerated performance (SSE4.1, AVX2, AVX-512 BITALG, NEON).
6969

7070
Installation
7171
-------------
@@ -182,6 +182,62 @@ immutable that is a very slow operation. Use a ``bytearray`` instead, and cast i
182182
Benchmark
183183
---------
184184

185+
For repeatable AVX2 and AVX-512 investigations, run the same checkout three
186+
times on each representative x86 machine:
187+
188+
.. code-block:: bash
189+
190+
scripts/benchmark_x86.sh before
191+
# Apply the candidate optimization, then:
192+
scripts/benchmark_x86.sh after
193+
194+
The script records CPU features and tool versions alongside Criterion output
195+
and end-to-end Python benchmark JSON. Compare results only between runs from
196+
the same machine.
197+
198+
AVX-512 results
199+
~~~~~~~~~~~~~~~~
200+
201+
Three-run medians on a Google Cloud ``c4-standard-4`` with an Intel Xeon
202+
Platinum 8581C (Emerald Rapids):
203+
204+
.. list-table::
205+
:header-rows: 1
206+
207+
* - Workload
208+
- Before
209+
- After
210+
- Speedup
211+
* - Python 1024x16 first, random/no-match
212+
- 3.222 us
213+
- 0.679 us
214+
- 4.75x
215+
* - Python 1024x16 best, random/no-match
216+
- 3.720 us
217+
- 0.651 us
218+
- 5.72x
219+
* - Python 1024x16 all, random/no-match
220+
- 3.466 us
221+
- 0.710 us
222+
- 4.88x
223+
* - Python 1024x32 first, random/no-match
224+
- 3.199 us
225+
- 1.278 us
226+
- 2.50x
227+
* - Python 1024x32 best, random/no-match
228+
- 3.729 us
229+
- 1.383 us
230+
- 2.70x
231+
* - Python 1024x32 all, random/no-match
232+
- 3.445 us
233+
- 1.377 us
234+
- 2.50x
235+
236+
The AVX-512 byte kernel also uses masked loads below 64 bytes, improving the
237+
measured 16-, 32-, 48-, and 63-byte Rust paths by 33%, 50%, 70%, and 194%
238+
respectively. AVX2-only tuning remains hardware-dependent and should be
239+
measured separately on a machine without AVX-512.
240+
185241
All benchmarks were run on an Apple M4 Max (ARM64, 16 logical cores, 64 GiB)
186242
with hexhamming v3.0.0, ``rustc`` 1.97.1, and Python 3.14.6. Values are the
187243
median of the means from three independent runs.

benches/hamming_bench.rs

Lines changed: 93 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use hexhamming::hex_hamming_distance_pack;
99

1010
// Hex sizes are character counts; byte sizes are the corresponding decoded lengths.
1111
const HEX_SIZES: [usize; 5] = [16, 32, 64, 128, 254];
12-
const BYTE_SIZES: [usize; 8] = [8, 16, 32, 64, 127, 128, 256, 512];
12+
const BYTE_SIZES: [usize; 12] = [8, 16, 32, 48, 63, 64, 65, 96, 127, 128, 256, 512];
1313

1414
fn pseudo_random_bytes(len: usize, seed: u64) -> Vec<u8> {
1515
let mut state = seed;
@@ -58,7 +58,7 @@ fn bench_hex_by_algo(c: &mut Criterion) {
5858
/// Benchmark bytes hamming distance across all available algorithms
5959
fn bench_bytes_by_algo(c: &mut Criterion) {
6060
let algos: &[&str] = if cfg!(target_arch = "x86_64") {
61-
&["classic", "sse", "avx2", "avx512"]
61+
&["classic", "native", "sse", "avx2", "avx512"]
6262
} else if cfg!(target_arch = "aarch64") {
6363
&["classic", "native", "neon"]
6464
} else {
@@ -358,6 +358,83 @@ fn bench_fixed_width_array_matrix(c: &mut Criterion) {
358358
}
359359
}
360360

361+
/// A/B comparison between the algorithm paths that DO and DO NOT engage the
362+
/// fixed-width cross-record scanner. `select_array_scanner_for_width` opts in
363+
/// only for `native`/`neon` on aarch64 and `native`/`avx512` on x86; other
364+
/// algorithms fall back to the per-record byte kernel. Running the same
365+
/// scenarios under both toggles measures the end-to-end dispatch alternatives.
366+
///
367+
/// Only runs on architectures where a fixed-width scanner is available.
368+
#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
369+
fn bench_fixed_width_scanner_vs_kernel(c: &mut Criterion) {
370+
// The x86 comparison includes both cross-record batching and the wider
371+
// AVX-512 BITALG popcount; it is an end-to-end comparison against the
372+
// narrower AVX2 per-record fallback, not an isolated batching benchmark.
373+
let pairs: &[(&str, &str)] = if cfg!(target_arch = "x86_64") {
374+
&[("avx512", "avx2")]
375+
} else {
376+
&[("native", "classic")]
377+
};
378+
379+
for &(scanner_algo, kernel_algo) in pairs {
380+
let scanner_ok = set_algorithm(scanner_algo).is_ok();
381+
set_algorithm("native").ok();
382+
let kernel_ok = set_algorithm(kernel_algo).is_ok();
383+
set_algorithm("native").ok();
384+
if !scanner_ok || !kernel_ok {
385+
continue;
386+
}
387+
388+
for &(role, algo) in &[("scanner", scanner_algo), ("kernel", kernel_algo)] {
389+
if set_algorithm(algo).is_err() {
390+
continue;
391+
}
392+
let mut group = c.benchmark_group(format!(
393+
"array_scanner/{scanner_algo}_vs_{kernel_algo}/{role}"
394+
));
395+
for &width in &[16usize, 32] {
396+
// Same random-no-match scenario as the C4 baseline (see the
397+
// catalog benchmarks in `array_api/512x16_random_no_match`) to
398+
// allow before/after comparison at that data point.
399+
for &count in &[512usize, 1024] {
400+
let small = pseudo_random_bytes(width, 0x51 + width as u64);
401+
let big = pseudo_random_bytes(count * width, 0xA1 + width as u64);
402+
group.bench_function(format!("{width}byte/{count}/first"), |bencher| {
403+
bencher.iter(|| {
404+
bytes_array_first_within_dist(
405+
black_box(&big),
406+
black_box(&small),
407+
black_box(0),
408+
)
409+
})
410+
});
411+
group.bench_function(format!("{width}byte/{count}/best"), |bencher| {
412+
bencher.iter(|| {
413+
bytes_array_best_within_dist(
414+
black_box(&big),
415+
black_box(&small),
416+
black_box(0),
417+
)
418+
})
419+
});
420+
group.bench_function(format!("{width}byte/{count}/all"), |bencher| {
421+
bencher.iter(|| {
422+
bytes_array_all_within_dist(
423+
black_box(&big),
424+
black_box(&small),
425+
black_box(0),
426+
)
427+
})
428+
});
429+
}
430+
}
431+
group.finish();
432+
}
433+
}
434+
435+
set_algorithm("native").ok();
436+
}
437+
361438
fn bench_fixed_width_parallel_crossover(c: &mut Criterion) {
362439
const PAR_THRESHOLDS: [(&str, usize); 2] = [
363440
("legacy", 5 * 1024 * 1024),
@@ -428,10 +505,23 @@ criterion_group!(
428505
bench_array_api,
429506
bench_array_random_and_boundaries,
430507
bench_fixed_width_array_matrix,
508+
bench_fixed_width_scanner_vs_kernel,
431509
bench_fixed_width_parallel_crossover,
432510
bench_hex_string_pack
433511
);
434-
#[cfg(not(target_arch = "aarch64"))]
512+
#[cfg(target_arch = "x86_64")]
513+
criterion_group!(
514+
benches,
515+
bench_hex_by_algo,
516+
bench_bytes_by_algo,
517+
bench_bytes_within_dist,
518+
bench_array_api,
519+
bench_array_random_and_boundaries,
520+
bench_fixed_width_array_matrix,
521+
bench_fixed_width_scanner_vs_kernel,
522+
bench_fixed_width_parallel_crossover
523+
);
524+
#[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
435525
criterion_group!(
436526
benches,
437527
bench_hex_by_algo,

scripts/benchmark_x86.sh

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#!/usr/bin/env bash
2+
set -euo pipefail
3+
4+
if [[ "$(uname -m)" != "x86_64" ]]; then
5+
echo "error: x86 SIMD benchmarks require an x86_64 machine" >&2
6+
exit 1
7+
fi
8+
9+
label="${1:-working-tree}"
10+
output_dir="${2:-benchmark-results/${label}}"
11+
rust_filter="${3:-}"
12+
mkdir -p "${output_dir}"
13+
14+
{
15+
echo "label=${label}"
16+
echo "date=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
17+
echo "kernel=$(uname -a)"
18+
echo "rustc=$(rustc --version)"
19+
echo "cargo=$(cargo --version)"
20+
echo "python=$(python3 --version 2>&1)"
21+
if command -v lscpu >/dev/null 2>&1; then
22+
lscpu
23+
elif command -v sysctl >/dev/null 2>&1; then
24+
sysctl -a 2>/dev/null | grep -E 'machdep.cpu.(brand_string|features|leaf7_features)'
25+
fi
26+
} >"${output_dir}/system.txt"
27+
28+
python3 -m pip install --quiet pytest pytest-benchmark
29+
python3 -m pip install --quiet .
30+
31+
for run in 1 2 3; do
32+
echo "Rust benchmark run ${run}/3"
33+
if [[ -n "${rust_filter}" ]]; then
34+
cargo bench --bench hamming_bench -- \
35+
"${rust_filter}" \
36+
--noplot \
37+
--save-baseline "${label}-rust-${run}" \
38+
2>&1 | tee "${output_dir}/criterion-${run}.txt"
39+
else
40+
cargo bench --bench hamming_bench -- \
41+
--noplot \
42+
--save-baseline "${label}-rust-${run}" \
43+
2>&1 | tee "${output_dir}/criterion-${run}.txt"
44+
fi
45+
46+
if [[ "${SKIP_PYTHON:-0}" != "1" ]]; then
47+
echo "Python benchmark run ${run}/3"
48+
python3 -m pytest test/ \
49+
-k bench \
50+
--benchmark-only \
51+
--benchmark-disable-gc \
52+
--benchmark-json="${output_dir}/python-${run}.json"
53+
fi
54+
done
55+
56+
echo "Results written to ${output_dir}"

src/api.rs

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,38 @@ fn select_array_scanner_for_width(width: usize) -> Option<ArrayScanner> {
5454
};
5555
}
5656

57-
#[cfg(not(target_arch = "aarch64"))]
57+
#[cfg(target_arch = "x86_64")]
58+
{
59+
// Only opt in to the AVX-512 cross-record scanners when the user
60+
// hasn't explicitly requested a narrower or scalar backend, and only
61+
// when the host actually supports the required feature set. This
62+
// preserves observable behavior on hosts without AVX-512 BITALG and
63+
// for callers that explicitly set "classic"/"sse"/"avx2".
64+
let algo = CURRENT_ALGO.load(Ordering::Relaxed);
65+
if (algo == ALGO_NATIVE || algo == ALGO_AVX512)
66+
&& is_x86_feature_detected!("avx512f")
67+
&& is_x86_feature_detected!("avx512bw")
68+
&& is_x86_feature_detected!("avx512bitalg")
69+
{
70+
return match width {
71+
16 => Some(ArrayScanner {
72+
first: crate::x86_simd::array_first_avx512_16_dispatch,
73+
best: crate::x86_simd::array_best_avx512_16_dispatch,
74+
all: crate::x86_simd::array_all_avx512_16_dispatch,
75+
}),
76+
32 => Some(ArrayScanner {
77+
first: crate::x86_simd::array_first_avx512_32_dispatch,
78+
best: crate::x86_simd::array_best_avx512_32_dispatch,
79+
all: crate::x86_simd::array_all_avx512_32_dispatch,
80+
}),
81+
_ => None,
82+
};
83+
}
84+
let _ = width;
85+
None
86+
}
87+
88+
#[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5889
{
5990
let _ = width;
6091
None

src/tests.rs

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -511,8 +511,22 @@ fn array_oracle(
511511

512512
#[test]
513513
fn test_fixed_width_array_scanners_match_randomized_oracle() {
514-
for algorithm in ["native", "classic"] {
515-
crate::set_algorithm(algorithm).unwrap();
514+
// On x86 the AVX-512 cross-record scanners are only reachable when the
515+
// active algorithm is `native` or `avx512`, so iterate through both to
516+
// exercise the specialized (widths 16 & 32) and the generic paths. On
517+
// aarch64 the same iteration covers the NEON fixed-width scanners and
518+
// the scalar fallback. `set_algorithm` returns `Err` when the CPU lacks
519+
// the requested extension, and we silently skip that iteration so this
520+
// test remains portable.
521+
let algorithms: &[&str] = if cfg!(target_arch = "x86_64") {
522+
&["native", "avx512", "classic"]
523+
} else {
524+
&["native", "classic"]
525+
};
526+
for &algorithm in algorithms {
527+
if crate::set_algorithm(algorithm).is_err() {
528+
continue;
529+
}
516530
for &width in &[1usize, 3, 7, 15, 16, 17, 31, 32, 33] {
517531
let count = 37;
518532
let mut state = 0xA5A5_1234_5678_9ABCu64 ^ width as u64;

0 commit comments

Comments
 (0)