Speed up AVX-512 byte and fixed-width array scans

README.rst

@@ -65,7 +65,7 @@ Lastly, I wanted to minimize dependencies, meaning you do not need to install
 As of v3.0.0, ``hexhamming`` is written in Rust using `PyO3 <https://pyo3.rs>`_
 and `maturin <https://www.maturin.rs>`_, providing memory safety, GIL release
 during computation, and free-threaded Python support while maintaining the same
-SIMD-accelerated performance (SSE4.1, AVX2, NEON).
+SIMD-accelerated performance (SSE4.1, AVX2, AVX-512 BITALG, NEON).
 
 Installation
 -------------
@@ -182,6 +182,62 @@ immutable that is a very slow operation. Use a ``bytearray`` instead, and cast i
 Benchmark
 ---------
 
+For repeatable AVX2 and AVX-512 investigations, run the same checkout three
+times on each representative x86 machine:
+
+.. code-block:: bash
+
+    scripts/benchmark_x86.sh before
+    # Apply the candidate optimization, then:
+    scripts/benchmark_x86.sh after
+
+The script records CPU features and tool versions alongside Criterion output
+and end-to-end Python benchmark JSON. Compare results only between runs from
+the same machine.
+
+AVX-512 results
+~~~~~~~~~~~~~~~~
+
+Three-run medians on a Google Cloud ``c4-standard-4`` with an Intel Xeon
+Platinum 8581C (Emerald Rapids):
+
+.. list-table::
+   :header-rows: 1
+
+   * - Workload
+     - Before
+     - After
+     - Speedup
+   * - Python 1024x16 first, random/no-match
+     - 3.222 us
+     - 0.679 us
+     - 4.75x
+   * - Python 1024x16 best, random/no-match
+     - 3.720 us
+     - 0.651 us
+     - 5.72x
+   * - Python 1024x16 all, random/no-match
+     - 3.466 us
+     - 0.710 us
+     - 4.88x
+   * - Python 1024x32 first, random/no-match
+     - 3.199 us
+     - 1.278 us
+     - 2.50x
+   * - Python 1024x32 best, random/no-match
+     - 3.729 us
+     - 1.383 us
+     - 2.70x
+   * - Python 1024x32 all, random/no-match
+     - 3.445 us
+     - 1.377 us
+     - 2.50x
+
+The AVX-512 byte kernel also uses masked loads below 64 bytes, improving the
+measured 16-, 32-, 48-, and 63-byte Rust paths by 33%, 50%, 70%, and 194%
+respectively. AVX2-only tuning remains hardware-dependent and should be
+measured separately on a machine without AVX-512.
+
 All benchmarks were run on an Apple M4 Max (ARM64, 16 logical cores, 64 GiB)
 with hexhamming v3.0.0, ``rustc`` 1.97.1, and Python 3.14.6. Values are the
 median of the means from three independent runs.

benches/hamming_bench.rs

@@ -9,7 +9,7 @@ use hexhamming::hex_hamming_distance_pack;
 
 // Hex sizes are character counts; byte sizes are the corresponding decoded lengths.
 const HEX_SIZES: [usize; 5] = [16, 32, 64, 128, 254];
-const BYTE_SIZES: [usize; 8] = [8, 16, 32, 64, 127, 128, 256, 512];
+const BYTE_SIZES: [usize; 12] = [8, 16, 32, 48, 63, 64, 65, 96, 127, 128, 256, 512];
 
 fn pseudo_random_bytes(len: usize, seed: u64) -> Vec<u8> {
     let mut state = seed;
@@ -58,7 +58,7 @@ fn bench_hex_by_algo(c: &mut Criterion) {
 /// Benchmark bytes hamming distance across all available algorithms
 fn bench_bytes_by_algo(c: &mut Criterion) {
     let algos: &[&str] = if cfg!(target_arch = "x86_64") {
-        &["classic", "sse", "avx2", "avx512"]
+        &["classic", "native", "sse", "avx2", "avx512"]
     } else if cfg!(target_arch = "aarch64") {
         &["classic", "native", "neon"]
     } else {
@@ -358,6 +358,83 @@ fn bench_fixed_width_array_matrix(c: &mut Criterion) {
     }
 }
 
+/// A/B comparison between the algorithm paths that DO and DO NOT engage the
+/// fixed-width cross-record scanner. `select_array_scanner_for_width` opts in
+/// only for `native`/`neon` on aarch64 and `native`/`avx512` on x86; other
+/// algorithms fall back to the per-record byte kernel. Running the same
+/// scenarios under both toggles measures the end-to-end dispatch alternatives.
+///
+/// Only runs on architectures where a fixed-width scanner is available.
+#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
+fn bench_fixed_width_scanner_vs_kernel(c: &mut Criterion) {
+    // The x86 comparison includes both cross-record batching and the wider
+    // AVX-512 BITALG popcount; it is an end-to-end comparison against the
+    // narrower AVX2 per-record fallback, not an isolated batching benchmark.
+    let pairs: &[(&str, &str)] = if cfg!(target_arch = "x86_64") {
+        &[("avx512", "avx2")]
+    } else {
+        &[("native", "classic")]
+    };
+
+    for &(scanner_algo, kernel_algo) in pairs {
+        let scanner_ok = set_algorithm(scanner_algo).is_ok();
+        set_algorithm("native").ok();
+        let kernel_ok = set_algorithm(kernel_algo).is_ok();
+        set_algorithm("native").ok();
+        if !scanner_ok || !kernel_ok {
+            continue;
+        }
+
+        for &(role, algo) in &[("scanner", scanner_algo), ("kernel", kernel_algo)] {
+            if set_algorithm(algo).is_err() {
+                continue;
+            }
+            let mut group = c.benchmark_group(format!(
+                "array_scanner/{scanner_algo}_vs_{kernel_algo}/{role}"
+            ));
+            for &width in &[16usize, 32] {
+                // Same random-no-match scenario as the C4 baseline (see the
+                // catalog benchmarks in `array_api/512x16_random_no_match`) to
+                // allow before/after comparison at that data point.
+                for &count in &[512usize, 1024] {
+                    let small = pseudo_random_bytes(width, 0x51 + width as u64);
+                    let big = pseudo_random_bytes(count * width, 0xA1 + width as u64);
+                    group.bench_function(format!("{width}byte/{count}/first"), |bencher| {
+                        bencher.iter(|| {
+                            bytes_array_first_within_dist(
+                                black_box(&big),
+                                black_box(&small),
+                                black_box(0),
+                            )
+                        })
+                    });
+                    group.bench_function(format!("{width}byte/{count}/best"), |bencher| {
+                        bencher.iter(|| {
+                            bytes_array_best_within_dist(
+                                black_box(&big),
+                                black_box(&small),
+                                black_box(0),
+                            )
+                        })
+                    });
+                    group.bench_function(format!("{width}byte/{count}/all"), |bencher| {
+                        bencher.iter(|| {
+                            bytes_array_all_within_dist(
+                                black_box(&big),
+                                black_box(&small),
+                                black_box(0),
+                            )
+                        })
+                    });
+                }
+            }
+            group.finish();
+        }
+    }
+
+    set_algorithm("native").ok();
+}
+
 fn bench_fixed_width_parallel_crossover(c: &mut Criterion) {
     const PAR_THRESHOLDS: [(&str, usize); 2] = [
         ("legacy", 5 * 1024 * 1024),
@@ -428,10 +505,23 @@ criterion_group!(
     bench_array_api,
     bench_array_random_and_boundaries,
     bench_fixed_width_array_matrix,
+    bench_fixed_width_scanner_vs_kernel,
     bench_fixed_width_parallel_crossover,
     bench_hex_string_pack
 );
-#[cfg(not(target_arch = "aarch64"))]
+#[cfg(target_arch = "x86_64")]
+criterion_group!(
+    benches,
+    bench_hex_by_algo,
+    bench_bytes_by_algo,
+    bench_bytes_within_dist,
+    bench_array_api,
+    bench_array_random_and_boundaries,
+    bench_fixed_width_array_matrix,
+    bench_fixed_width_scanner_vs_kernel,
+    bench_fixed_width_parallel_crossover
+);
+#[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
 criterion_group!(
     benches,
     bench_hex_by_algo,

scripts/benchmark_x86.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ "$(uname -m)" != "x86_64" ]]; then
+    echo "error: x86 SIMD benchmarks require an x86_64 machine" >&2
+    exit 1
+fi
+
+label="${1:-working-tree}"
+output_dir="${2:-benchmark-results/${label}}"
+rust_filter="${3:-}"
+mkdir -p "${output_dir}"
+
+{
+    echo "label=${label}"
+    echo "date=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
+    echo "kernel=$(uname -a)"
+    echo "rustc=$(rustc --version)"
+    echo "cargo=$(cargo --version)"
+    echo "python=$(python3 --version 2>&1)"
+    if command -v lscpu >/dev/null 2>&1; then
+        lscpu
+    elif command -v sysctl >/dev/null 2>&1; then
+        sysctl -a 2>/dev/null | grep -E 'machdep.cpu.(brand_string|features|leaf7_features)'
+    fi
+} >"${output_dir}/system.txt"
+
+python3 -m pip install --quiet pytest pytest-benchmark
+python3 -m pip install --quiet .
+
+for run in 1 2 3; do
+    echo "Rust benchmark run ${run}/3"
+    if [[ -n "${rust_filter}" ]]; then
+        cargo bench --bench hamming_bench -- \
+            "${rust_filter}" \
+            --noplot \
+            --save-baseline "${label}-rust-${run}" \
+            2>&1 | tee "${output_dir}/criterion-${run}.txt"
+    else
+        cargo bench --bench hamming_bench -- \
+            --noplot \
+            --save-baseline "${label}-rust-${run}" \
+            2>&1 | tee "${output_dir}/criterion-${run}.txt"
+    fi
+
+    if [[ "${SKIP_PYTHON:-0}" != "1" ]]; then
+        echo "Python benchmark run ${run}/3"
+        python3 -m pytest test/ \
+            -k bench \
+            --benchmark-only \
+            --benchmark-disable-gc \
+            --benchmark-json="${output_dir}/python-${run}.json"
+    fi
+done
+
+echo "Results written to ${output_dir}"

src/api.rs

@@ -54,7 +54,38 @@ fn select_array_scanner_for_width(width: usize) -> Option<ArrayScanner> {
         };
     }
 
-    #[cfg(not(target_arch = "aarch64"))]
+    #[cfg(target_arch = "x86_64")]
+    {
+        // Only opt in to the AVX-512 cross-record scanners when the user
+        // hasn't explicitly requested a narrower or scalar backend, and only
+        // when the host actually supports the required feature set. This
+        // preserves observable behavior on hosts without AVX-512 BITALG and
+        // for callers that explicitly set "classic"/"sse"/"avx2".
+        let algo = CURRENT_ALGO.load(Ordering::Relaxed);
+        if (algo == ALGO_NATIVE || algo == ALGO_AVX512)
+            && is_x86_feature_detected!("avx512f")
+            && is_x86_feature_detected!("avx512bw")
+            && is_x86_feature_detected!("avx512bitalg")
+        {
+            return match width {
+                16 => Some(ArrayScanner {
+                    first: crate::x86_simd::array_first_avx512_16_dispatch,
+                    best: crate::x86_simd::array_best_avx512_16_dispatch,
+                    all: crate::x86_simd::array_all_avx512_16_dispatch,
+                }),
+                32 => Some(ArrayScanner {
+                    first: crate::x86_simd::array_first_avx512_32_dispatch,
+                    best: crate::x86_simd::array_best_avx512_32_dispatch,
+                    all: crate::x86_simd::array_all_avx512_32_dispatch,
+                }),
+                _ => None,
+            };
+        }
+        let _ = width;
+        None
+    }
+
+    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
     {
         let _ = width;
         None

src/tests.rs

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