Speed up ARM fixed-width array scans (#54)

Add NEON cross-record scanners for 16- and 32-byte records, expand semantic and benchmark coverage, and retune the parallel crossover based on three-run measurements.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 01071e57-12fd-4f42-bc67-e2b274d0b007

Author: mrecachinas

README.rst

@@ -183,7 +183,7 @@ Benchmark
 ---------
 
 All benchmarks were run on an Apple M4 Max (ARM64, 16 logical cores, 64 GiB)
-with hexhamming v3.0.0, ``rustc`` 1.96.1, and Python 3.14.6. Values are the
+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.
 
 Raw Rust (no Python overhead)
@@ -192,6 +192,24 @@ Raw Rust (no Python overhead)
 These numbers show the pure computation time using Rust's ``criterion`` benchmarks
 (``cargo bench --no-default-features``), with no Python/PyO3 overhead.
 
+Issue #51 fixed-width array matrix (1024 records; median of three run medians)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The matrix uses deterministic random no-match and exact-midpoint cases for
+16-byte and 32-byte records. Each run uses
+``--warm-up-time 1 --measurement-time 1 --sample-size 20``.
+
+====================================  ===========  ===========
+Case                                 16-byte (ns) 32-byte (ns)
+====================================  ===========  ===========
+random no-match / first                    397.5        797.8
+random no-match / best                     407.0        992.2
+random no-match / all                      523.7       1047.7
+exact midpoint / first                    216.3        422.7
+exact midpoint / best                     216.6        414.0
+exact midpoint / all                      540.9        822.0
+====================================  ===========  ===========
+
 ================================================  ===========
 Name                                              Mean (ns)
 ================================================  ===========
@@ -257,6 +275,23 @@ all_within_dist [512×16, at end]                             537.5
 all_within_dist [16384×64, mid]                           30,725.4
 ======================================================  ===========
 
+Issue #51 Python buffer matrix (1024 records; median of three run medians)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+These end-to-end timings use ``timeit.repeat`` with 10,000 calls per sample,
+including the PyO3 wrapper and buffer-protocol path.
+
+====================================  ===========  ===========
+Case                                 16-byte (ns) 32-byte (ns)
+====================================  ===========  ===========
+random no-match / first                    443.4        836.8
+random no-match / best                     487.0        850.2
+random no-match / all                      564.8        851.7
+exact midpoint / first                    272.2        468.1
+exact midpoint / best                     295.0        483.8
+exact midpoint / all                      638.8        921.0
+====================================  ===========  ===========
+
 For random inputs, the direct APIs also avoid the temporary big integers used
 by an equivalent standard-library implementation:
 
@@ -276,6 +311,7 @@ dominates (roughly 30–40 ns on this machine). For large inputs
 (1024+ chars, 16384-element arrays), computation dominates and Python overhead
 is negligible. Byte operations release the GIL at 16 KiB, while immutable
 strings use a zero-copy detached path from 4 KiB. Array wrappers release the GIL
-at 64 KiB and parallelize with Rayon at 5 MiB; the ``first`` variant additionally
-short-circuits on the first hit, so a match near the start is much faster than
-one near the end.
+at 64 KiB; generic scans parallelize with Rayon at 5 MiB, while the optimized
+16/32-byte NEON scanners use a measured 16 MiB crossover. The ``first`` variant
+additionally short-circuits on the first hit, so a match near the start is much
+faster than one near the end.

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; 5] = [8, 16, 32, 64, 127];
+const BYTE_SIZES: [usize; 8] = [8, 16, 32, 64, 127, 128, 256, 512];
 
 fn pseudo_random_bytes(len: usize, seed: u64) -> Vec<u8> {
     let mut state = seed;
@@ -287,6 +287,124 @@ fn bench_array_random_and_boundaries(c: &mut Criterion) {
     group.finish();
 }
 
+fn fixed_width_array_case(width: usize, scenario: &str) -> (Vec<u8>, Vec<u8>, i64) {
+    const NUM_ELEMENTS: usize = 1024;
+    let small = pseudo_random_bytes(width, 0x51 + width as u64);
+    let mut big = pseudo_random_bytes(NUM_ELEMENTS * width, 0xA1 + width as u64);
+    let (match_index, max_dist) = match scenario {
+        "random_no_match" => (None, 0),
+        "exact_early" => (Some(0), 0),
+        "exact_mid" => (Some(NUM_ELEMENTS / 2), 0),
+        "exact_late" => (Some(NUM_ELEMENTS - 1), 0),
+        "threshold_d_minus_1" => (Some(NUM_ELEMENTS / 2), 3),
+        "threshold_d" => (Some(NUM_ELEMENTS / 2), 4),
+        "threshold_d_plus_1" => (Some(NUM_ELEMENTS / 2), 5),
+        _ => unreachable!("unknown fixed-width benchmark scenario"),
+    };
+    if let Some(index) = match_index {
+        let record = &mut big[index * width..(index + 1) * width];
+        record.copy_from_slice(&small);
+        if scenario.starts_with("threshold_") {
+            record[0] ^= 0x0F;
+        }
+    }
+    (big, small, max_dist)
+}
+
+fn bench_fixed_width_array_matrix(c: &mut Criterion) {
+    let scenarios = [
+        "random_no_match",
+        "exact_early",
+        "exact_mid",
+        "exact_late",
+        "threshold_d_minus_1",
+        "threshold_d",
+        "threshold_d_plus_1",
+    ];
+
+    for width in [16usize, 32] {
+        let mut group = c.benchmark_group(format!("array_matrix/{width}byte_records"));
+        for scenario in scenarios {
+            let (big, small, max_dist) = fixed_width_array_case(width, scenario);
+            group.bench_function(format!("{scenario}/first"), |bencher| {
+                bencher.iter(|| {
+                    bytes_array_first_within_dist(
+                        black_box(&big),
+                        black_box(&small),
+                        black_box(max_dist),
+                    )
+                })
+            });
+            group.bench_function(format!("{scenario}/best"), |bencher| {
+                bencher.iter(|| {
+                    bytes_array_best_within_dist(
+                        black_box(&big),
+                        black_box(&small),
+                        black_box(max_dist),
+                    )
+                })
+            });
+            group.bench_function(format!("{scenario}/all"), |bencher| {
+                bencher.iter(|| {
+                    bytes_array_all_within_dist(
+                        black_box(&big),
+                        black_box(&small),
+                        black_box(max_dist),
+                    )
+                })
+            });
+        }
+        group.finish();
+    }
+}
+
+fn bench_fixed_width_parallel_crossover(c: &mut Criterion) {
+    const PAR_THRESHOLDS: [(&str, usize); 2] = [
+        ("legacy", 5 * 1024 * 1024),
+        ("fixed_width", 16 * 1024 * 1024),
+    ];
+    let mut group = c.benchmark_group("array_matrix/parallel_crossover");
+    group.sample_size(10);
+    for width in [16usize, 32] {
+        let small = pseudo_random_bytes(width, 0xC1 + width as u64);
+        for &(threshold_name, threshold_bytes) in &PAR_THRESHOLDS {
+            let threshold_elements = threshold_bytes / width;
+            for count in [
+                threshold_elements - 1,
+                threshold_elements,
+                threshold_elements + 1,
+            ] {
+                let big = pseudo_random_bytes(count * width, 0xD1 + count as u64);
+                group.bench_function(
+                    format!("{threshold_name}/{width}byte/{count}elements/best"),
+                    |bencher| {
+                        bencher.iter(|| {
+                            bytes_array_best_within_dist(
+                                black_box(&big),
+                                black_box(&small),
+                                black_box(0),
+                            )
+                        })
+                    },
+                );
+                group.bench_function(
+                    format!("{threshold_name}/{width}byte/{count}elements/all"),
+                    |bencher| {
+                        bencher.iter(|| {
+                            bytes_array_all_within_dist(
+                                black_box(&big),
+                                black_box(&small),
+                                black_box(0),
+                            )
+                        })
+                    },
+                );
+            }
+        }
+    }
+    group.finish();
+}
+
 #[cfg(target_arch = "aarch64")]
 fn bench_hex_string_pack(c: &mut Criterion) {
     // AArch64-only group for the packed NEON hex-string path.
@@ -309,6 +427,8 @@ criterion_group!(
     bench_bytes_within_dist,
     bench_array_api,
     bench_array_random_and_boundaries,
+    bench_fixed_width_array_matrix,
+    bench_fixed_width_parallel_crossover,
     bench_hex_string_pack
 );
 #[cfg(not(target_arch = "aarch64"))]
@@ -318,6 +438,8 @@ criterion_group!(
     bench_bytes_by_algo,
     bench_bytes_within_dist,
     bench_array_api,
-    bench_array_random_and_boundaries
+    bench_array_random_and_boundaries,
+    bench_fixed_width_array_matrix,
+    bench_fixed_width_parallel_crossover
 );
 criterion_main!(benches);

src/api.rs

@@ -1,3 +1,5 @@
+#[cfg(target_arch = "aarch64")]
+use crate::ALGO_NEON;
 use crate::{
     hamming_distance_bytes_dispatch, hamming_distance_string_dispatch,
     select_bytes_kernel_for_width, BytesKernel, ALGO_CLASSIC, ALGO_NATIVE, CURRENT_ALGO,
@@ -10,11 +12,55 @@ use std::sync::atomic::Ordering;
 
 /// Minimum total byte size of big_array before we use rayon parallel paths.
 const PAR_THRESHOLD_BYTES: usize = 5 * 1024 * 1024;
+/// The fixed-width NEON scanners make serial scans substantially cheaper, so
+/// use a larger crossover before paying Rayon scheduling and partition costs.
+const FIXED_WIDTH_PAR_THRESHOLD_BYTES: usize = 16 * 1024 * 1024;
 /// Keep byte-array scans to a small number of coarse jobs. More workers spend
 /// more time scheduling these very small per-record calculations than running
 /// them on current many-core CPUs.
 const PAR_JOBS: usize = 4;
 
+type ArrayFirstScanner = fn(&[u8], &[u8], i64) -> Option<usize>;
+type ArrayBestScanner = fn(&[u8], &[u8], i64) -> Option<(u64, usize)>;
+type ArrayAllScanner = fn(&[u8], &[u8], i64) -> Vec<(u64, usize)>;
+
+#[derive(Clone, Copy)]
+struct ArrayScanner {
+    first: ArrayFirstScanner,
+    best: ArrayBestScanner,
+    all: ArrayAllScanner,
+}
+
+#[inline]
+fn select_array_scanner_for_width(width: usize) -> Option<ArrayScanner> {
+    #[cfg(target_arch = "aarch64")]
+    {
+        let algo = CURRENT_ALGO.load(Ordering::Relaxed);
+        if algo != ALGO_NATIVE && algo != ALGO_NEON {
+            return None;
+        }
+        return match width {
+            16 => Some(ArrayScanner {
+                first: crate::neon_simd::array_first_neon_16,
+                best: crate::neon_simd::array_best_neon_16,
+                all: crate::neon_simd::array_all_neon_16,
+            }),
+            32 => Some(ArrayScanner {
+                first: crate::neon_simd::array_first_neon_32,
+                best: crate::neon_simd::array_best_neon_32,
+                all: crate::neon_simd::array_all_neon_32,
+            }),
+            _ => None,
+        };
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        let _ = width;
+        None
+    }
+}
+
 #[inline]
 fn partition_element_ranges(num_elements: usize) -> [(usize, usize); PAR_JOBS] {
     let base = num_elements / PAR_JOBS;
@@ -103,6 +149,9 @@ pub fn bytes_array_first_within_dist(
     if big_array.len() % small_array.len() != 0 {
         return Err("array_of_elems size must be multiplier of elem_to_compare");
     }
+    if let Some(scanner) = select_array_scanner_for_width(small_array.len()) {
+        return Ok((scanner.first)(big_array, small_array, max_dist));
+    }
     // `first` has early-exit semantics: the serial scan returns as soon as the
     // first match is found, which is essentially free for early/common matches.
     // Parallelizing this requires a full non-short-circuiting scan to compute
@@ -149,7 +198,16 @@ pub fn bytes_array_best_within_dist(
         return Err("array_of_elems size must be multiplier of elem_to_compare");
     }
     let kernel = select_bytes_kernel_for_width(small_array.len());
-    if big_array.len() < PAR_THRESHOLD_BYTES {
+    let scanner = select_array_scanner_for_width(small_array.len());
+    let parallel_threshold = if scanner.is_some() {
+        FIXED_WIDTH_PAR_THRESHOLD_BYTES
+    } else {
+        PAR_THRESHOLD_BYTES
+    };
+    if big_array.len() < parallel_threshold {
+        if let Some(scanner) = scanner {
+            return Ok((scanner.best)(big_array, small_array, max_dist));
+        }
         return Ok(serial_best_within_dist(
             big_array,
             small_array,
@@ -166,8 +224,11 @@ pub fn bytes_array_best_within_dist(
         .with_max_len(1)
         .map(|&(start, end)| {
             let chunk = &big_array[start * elem_size..end * elem_size];
-            serial_best_within_dist(chunk, small_array, max_dist, kernel)
-                .map(|(distance, index)| (distance, index + start))
+            let best = match scanner {
+                Some(scanner) => (scanner.best)(chunk, small_array, max_dist),
+                None => serial_best_within_dist(chunk, small_array, max_dist, kernel),
+            };
+            best.map(|(distance, index)| (distance, index + start))
         })
         .reduce(|| None, merge_best))
 }
@@ -225,7 +286,16 @@ pub fn bytes_array_all_within_dist(
         return Err("array_of_elems size must be multiplier of elem_to_compare");
     }
     let kernel = select_bytes_kernel_for_width(small_array.len());
-    if big_array.len() < PAR_THRESHOLD_BYTES {
+    let scanner = select_array_scanner_for_width(small_array.len());
+    let parallel_threshold = if scanner.is_some() {
+        FIXED_WIDTH_PAR_THRESHOLD_BYTES
+    } else {
+        PAR_THRESHOLD_BYTES
+    };
+    if big_array.len() < parallel_threshold {
+        if let Some(scanner) = scanner {
+            return Ok((scanner.all)(big_array, small_array, max_dist));
+        }
         return Ok(serial_all_within_dist(
             big_array,
             small_array,
@@ -241,7 +311,11 @@ pub fn bytes_array_all_within_dist(
         .with_max_len(1)
         .map(|&(start, end)| {
             let chunk = &big_array[start * elem_size..end * elem_size];
-            serial_all_within_dist(chunk, small_array, max_dist, kernel)
+            let matches = match scanner {
+                Some(scanner) => (scanner.all)(chunk, small_array, max_dist),
+                None => serial_all_within_dist(chunk, small_array, max_dist, kernel),
+            };
+            matches
                 .into_iter()
                 .map(|(distance, index)| (distance, index + start))
                 .collect()