Skip to content

perf: optimize FFI boundary and SIMD hot paths; fix AVX-512 string overflow #48

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mrecachinas merged 3 commits into from
May 29, 2026
Merged

perf: optimize FFI boundary and SIMD hot paths; fix AVX-512 string overflow #48

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
91d5337
perf: optimize FFI boundary and SIMD hot paths; fix AVX-512 string ov…
mrecachinas May 29, 2026
7a7be92
perf: keep first_within_dist serial to preserve early-exit short-circuit
mrecachinas May 29, 2026
0369c1c
perf: mark serial batch helpers #[inline]
mrecachinas May 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 11 additions & 13 deletions src/api.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,19 +82,15 @@ 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 big_array.len() < PAR_THRESHOLD_BYTES {
return Ok(serial_first_within_dist(big_array, small_array, max_dist));
}
let elem_size = small_array.len();
Ok(big_array
.par_chunks_exact(elem_size)
.enumerate()
.filter_map(|(i, chunk)| {
(hamming_distance_bytes_dispatch(chunk, small_array, max_dist) != u64::MAX).then_some(i)
})
.min())
// `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
// the minimum matching index, which is dramatically slower for early/mid
// matches and cannot beat serial for a match at index 0. Always go serial.
Ok(serial_first_within_dist(big_array, small_array, max_dist))
}

#[inline]
fn serial_first_within_dist(big_array: &[u8], small_array: &[u8], max_dist: i64) -> Option<usize> {
let elem_size = small_array.len();
let num_elements = big_array.len() / elem_size;
Expand Down Expand Up @@ -164,6 +160,7 @@ pub fn bytes_array_best_within_dist(
))
}

#[inline]
fn serial_best_within_dist(
big_array: &[u8],
small_array: &[u8],
Expand Down Expand Up @@ -222,6 +219,7 @@ pub fn bytes_array_all_within_dist(
Ok(results)
}

#[inline]
fn serial_all_within_dist(
big_array: &[u8],
small_array: &[u8],
Expand Down Expand Up @@ -440,7 +438,7 @@ mod tests {

#[test]
fn first_within_dist_large_batch() {
// Above PAR_THRESHOLD_BYTES → parallel path
// Large batch → serial early-exit path (first is never parallelized)
let elem_size = 16;
let n = 100_000; // 1.6 MB > 64KB
let needle = vec![0x00u8; elem_size];
Expand All @@ -452,7 +450,7 @@ mod tests {
}

#[test]
fn first_within_dist_parallel_returns_lowest_index() {
fn first_within_dist_returns_lowest_index() {
let elem_size = 16;
let n = 100_000;
let needle = vec![0x00u8; elem_size];
Expand Down
231 changes: 110 additions & 121 deletions src/neon_simd.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,46 @@ unsafe fn hex_parse_neon(
let bad_letter = vandq_u8(is_letter, vcltq_u8(adjusted, ten));
vorrq_u8(result, bad_letter)
}

/// Parse 32 hex chars at `a`/`b`, XOR their nibbles, and pack the 32 nibble-XOR
/// results into 16 bytes (each byte holds two XOR'd nibbles).
///
/// Returns `(packed, bad)` where `packed` is ready for `vcntq_u8` popcount and
/// `bad` is a per-lane mask (non-zero lane ⇒ an invalid hex char was seen).
/// The caller is responsible for accumulating `bad` and validating once.
///
/// SAFETY: `a` and `b` must each be valid for 32 readable bytes.
#[inline(always)]
unsafe fn pack32_xor_neon(
a: *const u8,
b: *const u8,
case_mask: uint8x16_t,
ascii_0: uint8x16_t,
seven: uint8x16_t,
nine: uint8x16_t,
ten: uint8x16_t,
fifteen_u: uint8x16_t,
) -> (uint8x16_t, uint8x16_t) {
let a_lo = hex_parse_neon(vld1q_u8(a), case_mask, ascii_0, seven, nine, ten);
let b_lo = hex_parse_neon(vld1q_u8(b), case_mask, ascii_0, seven, nine, ten);
let a_hi = hex_parse_neon(vld1q_u8(a.add(16)), case_mask, ascii_0, seven, nine, ten);
let b_hi = hex_parse_neon(vld1q_u8(b.add(16)), case_mask, ascii_0, seven, nine, ten);

// Per-lane invalid mask — caller accumulates and checks once.
let bad = vorrq_u8(
vcgtq_u8(vorrq_u8(a_lo, b_lo), fifteen_u),
vcgtq_u8(vorrq_u8(a_hi, b_hi), fifteen_u),
);

// XOR nibbles, then pack even/odd nibbles into bytes.
let xor_lo = veorq_u8(a_lo, b_lo);
let xor_hi = veorq_u8(a_hi, b_hi);
let even = vuzp1q_u8(xor_lo, xor_hi);
let odd = vuzp2q_u8(xor_lo, xor_hi);
// Constant left-shift by 4 (nibbles are 0-15, so logical shift is correct).
let packed = vorrq_u8(vshlq_n_u8(even, 4), odd);
(packed, bad)
}
/// Alternative: parse 32 hex chars → pack into 16 bytes → use vcntq_u8.
/// Processes 32 hex chars per iteration (vs 16 in the nibble approach),
/// and replaces the vqtbl1q popcount lookup with the native vcntq_u8
Expand All @@ -312,83 +352,62 @@ pub unsafe fn hamming_distance_string_neon_pack(a: &[u8], b: &[u8]) -> Result<u6
let seven = vdupq_n_u8(7);
let nine = vdupq_n_u8(9);
let ten = vdupq_n_u8(10);
let four = vdupq_n_s8(4);
let zero = vdupq_n_u8(0);

let mut i = 0usize;
let mut difference: u64 = 0;
// Accumulate invalid-char masks across the whole SIMD region and check once.
let mut bad_acc = zero;

// Each 32-char iteration popcounts into a u8 lane (≤8 per packed byte), so
// up to 31 iterations are safe before a lane overflows (31*8=248 < 256).
// Batch BATCH iterations per horizontal reduction to keep the hot loop free
// of cross-lane reductions.
const BATCH: usize = 16;

// Process 32 hex chars at a time:
// Load 32 chars (2×16), parse to nibbles, pack pairs into bytes,
// XOR the packed bytes, vcntq_u8 popcount.
while i + 32 * BATCH <= length {
let mut acc = zero;
for _ in 0..BATCH {
let (packed, bad) = pack32_xor_neon(
a.as_ptr().add(i),
b.as_ptr().add(i),
case_mask,
ascii_0,
seven,
nine,
ten,
fifteen_u,
);
bad_acc = vorrq_u8(bad_acc, bad);
acc = vaddq_u8(acc, vcntq_u8(packed));
i += 32;
}
difference += vaddlvq_u8(acc) as u64;
}

// Remaining 32-char chunks (< BATCH of them — still safe in u8 lanes).
let mut acc = zero;
while i + 32 <= length {
// Parse first 16 hex chars → nibbles
let a_lo = hex_parse_neon(
vld1q_u8(a.as_ptr().add(i)),
let (packed, bad) = pack32_xor_neon(
a.as_ptr().add(i),
b.as_ptr().add(i),
case_mask,
ascii_0,
seven,
nine,
ten,
fifteen_u,
);
let b_lo = hex_parse_neon(
vld1q_u8(b.as_ptr().add(i)),
case_mask,
ascii_0,
seven,
nine,
ten,
);
// Parse next 16 hex chars → nibbles
let a_hi = hex_parse_neon(
vld1q_u8(a.as_ptr().add(i + 16)),
case_mask,
ascii_0,
seven,
nine,
ten,
);
let b_hi = hex_parse_neon(
vld1q_u8(b.as_ptr().add(i + 16)),
case_mask,
ascii_0,
seven,
nine,
ten,
);

// Validate all 4 vectors — single cmpgt(or(a|b), 15)
let bad = vorrq_u8(
vcgtq_u8(vorrq_u8(a_lo, b_lo), fifteen_u),
vcgtq_u8(vorrq_u8(a_hi, b_hi), fifteen_u),
);
if vmaxvq_u8(bad) != 0 {
return Err("hex string contains invalid char");
}

// XOR nibbles first (before packing — same result, fewer packs)
let xor_lo = veorq_u8(a_lo, b_lo); // 16 nibble XOR results
let xor_hi = veorq_u8(a_hi, b_hi); // 16 nibble XOR results

// Pack: interleave even/odd nibbles into bytes
// xor_lo has nibbles [0,1,2,3,...,15], xor_hi has [16,17,...,31]
// We want bytes where each byte = (nibble[2k] << 4) | nibble[2k+1]
// Use UZP to deinterleave even/odd lanes, then shift+OR
let even_lo = vuzp1q_u8(xor_lo, xor_hi); // even indices: 0,2,4,...
let odd_lo = vuzp2q_u8(xor_lo, xor_hi); // odd indices: 1,3,5,...
// Shift even nibbles left 4 bits via reinterpret + signed shift
let packed = vorrq_u8(
vreinterpretq_u8_s8(vshlq_s8(vreinterpretq_s8_u8(even_lo), four)),
odd_lo,
);

// Now packed has 16 bytes, each containing two XOR'd nibbles
// vcntq_u8 counts all set bits per byte — exactly what we want
let cnt = vcntq_u8(packed);
// Horizontal sum
difference += vaddlvq_u8(cnt) as u64;

bad_acc = vorrq_u8(bad_acc, bad);
acc = vaddq_u8(acc, vcntq_u8(packed));
i += 32;
}
difference += vaddlvq_u8(acc) as u64;

// Single validation for the entire 32-char SIMD region.
if vmaxvq_u8(bad_acc) != 0 {
return Err("hex string contains invalid char");
}

// Handle remaining chars with the nibble-based approach
// §13: popcnt_tbl loaded once at function entry scope, not per iteration
Expand Down Expand Up @@ -456,73 +475,43 @@ pub unsafe fn hamming_distance_string_neon_pack_with_max(
let seven = vdupq_n_u8(7);
let nine = vdupq_n_u8(9);
let ten = vdupq_n_u8(10);
let four = vdupq_n_s8(4);
let zero = vdupq_n_u8(0);

let mut i = 0usize;
let mut difference: u64 = 0;
let mut bad_acc = zero;

// Process 32 hex chars at a time with threshold check after each iteration
// Process 32 hex chars at a time, batching cross-lane reductions and the
// threshold check. Each packed byte popcount is ≤8, so BATCH iterations are
// safe in u8 lanes (BATCH*8 < 256). Early exit is granular to one batch.
const BATCH: usize = 16;
while i + 32 <= length {
let a_lo = hex_parse_neon(
vld1q_u8(a.as_ptr().add(i)),
case_mask,
ascii_0,
seven,
nine,
ten,
);
let b_lo = hex_parse_neon(
vld1q_u8(b.as_ptr().add(i)),
case_mask,
ascii_0,
seven,
nine,
ten,
);
let a_hi = hex_parse_neon(
vld1q_u8(a.as_ptr().add(i + 16)),
case_mask,
ascii_0,
seven,
nine,
ten,
);
let b_hi = hex_parse_neon(
vld1q_u8(b.as_ptr().add(i + 16)),
case_mask,
ascii_0,
seven,
nine,
ten,
);

// Validate
let bad = vorrq_u8(
vcgtq_u8(vorrq_u8(a_lo, b_lo), fifteen_u),
vcgtq_u8(vorrq_u8(a_hi, b_hi), fifteen_u),
);
if vmaxvq_u8(bad) != 0 {
let mut acc = zero;
let mut n = 0;
while n < BATCH && i + 32 <= length {
let (packed, bad) = pack32_xor_neon(
a.as_ptr().add(i),
b.as_ptr().add(i),
case_mask,
ascii_0,
seven,
nine,
ten,
fifteen_u,
);
bad_acc = vorrq_u8(bad_acc, bad);
acc = vaddq_u8(acc, vcntq_u8(packed));
i += 32;
n += 1;
}
// Invalid hex chars take precedence over the max_dist sentinel.
if vmaxvq_u8(bad_acc) != 0 {
return Err("hex string contains invalid char");
}

let xor_lo = veorq_u8(a_lo, b_lo);
let xor_hi = veorq_u8(a_hi, b_hi);

let even_lo = vuzp1q_u8(xor_lo, xor_hi);
let odd_lo = vuzp2q_u8(xor_lo, xor_hi);
let packed = vorrq_u8(
vreinterpretq_u8_s8(vshlq_s8(vreinterpretq_s8_u8(even_lo), four)),
odd_lo,
);

let cnt = vcntq_u8(packed);
difference += vaddlvq_u8(cnt) as u64;

difference += vaddlvq_u8(acc) as u64;
if difference > max_dist {
return Ok(u64::MAX);
}

i += 32;
}

// Handle remaining chars with the nibble-based approach
Expand Down
Loading
Loading