diff --git a/.claude/PERFORMANCE_LOG.md b/.claude/PERFORMANCE_LOG.md index 49bbc26..c0ae194 100644 --- a/.claude/PERFORMANCE_LOG.md +++ b/.claude/PERFORMANCE_LOG.md @@ -432,3 +432,48 @@ the instructions. The EPYC 7713 quick run showed the same pattern (indel x86 now uses the blockwise 64-bit scalar extend, as on ARM. Output is unchanged (same offsets). To confirm with `bench/sweep.py` on x86. + +- Confirmed on the EPYC 7713 (`bench/sweep.py`, `8ccfb64`, 3 runs per case, + vs the pre-fix quick run): `l10k_5` indel score 1.69× → 0.80× C, + `m1k_20` indel CIGAR 1.25× → 0.87×; edit unchanged (it rarely hit the + kernel). With #45, 1k/10k med/low went from a median of 0.94× to 0.45× + generic C. + +## 2026-10-15: BiWFA overlap scan vectorized (HPCBio/wfa2lib-rs#31) + +The full EPYC sweep after #45/#48 left affine2p BiWFA (ultralow) as the main +case slower than C: 1.08–1.15× generic C, ~2.2× C-native, at 1.2–1.4× C's +instructions. A `sample` profile on the M2 (`m1k_5` affine2p ultralow) put +73% of the time in `bialign_overlap`, against 11% each for compute and +extend. + +- At every step BiWFA checks the new wavefronts against the last + `max_score_scope` scores of the other direction: one M scan plus four + I1/D1/I2/D2 scans for affine2p (scope ~25 with the default penalties). + Rust had NEON only for the M scan; the indel scans and every scan on x86 + were scalar. C has AVX2/AVX-512 versions of both, used when built with + `-march=native`: the C-native gap. +- Fix: one `scan_overlap` for all five scans. Chunks of 8 diagonals are + tested at once (NEON: 2 × 4 lanes; AVX2: one `vpermd`-reversed load, + inlined with `target-cpu=native`, runtime-detected otherwise); only a chunk + with a candidate is rechecked one diagonal at a time. The visiting order is + unchanged, so the breakpoint is the same as the scalar scan (and C's). +- M2 (`data-quick`, Rust vs `main` vs C 3282f83, ultralow, CIGAR output): + +| Case | `main` | this | C | this / C | +|---|---|---|---|---| +| m1k_5 affine2p | 1612 ms | 732 ms | 1.58 s | 0.46 | +| m1k_5 affine2p m-1 | 7118 ms | 2799 ms | 7.45 s | 0.38 | +| l10k_5 affine2p | 1325 ms | 702 ms | 1.36 s | 0.52 | +| l10k_5 affine2p m-1 | 3576 ms | 1474 ms | 3.63 s | 0.41 | +| m1k_5 affine | 132 ms | 122 ms | 149 ms | 0.82 | +| m1k_5 edit / indel / linear | | unchanged | | 0.83 / 0.85 / 0.81 | + +- Output identical to C and `main` in all cases above, and in x86 builds + (AVX2 and generic, run under Rosetta) on subsets. Differential harness: + 1242/1262 identical, the 20 known #22 cases. New unit tests compare the + scan with a naive scan (tails, nulls, rejected indel hits) and the SIMD + chunk tests with scalar; both catch a deliberately broken lane pairing + and a scan that skips the rest of a chunk. +- To confirm on x86 with `bench/sweep.py` (expect affine2p ultralow near or + below C-native). diff --git a/.claude/PROGRESS.md b/.claude/PROGRESS.md index d15899a..c7fe60d 100644 --- a/.claude/PROGRESS.md +++ b/.claude/PROGRESS.md @@ -168,4 +168,8 @@ reports, inbound C ports) is tracked in HPCBio/wfa2lib-rs#19. up to 2.16× slower than C on long/divergent edit/indel/linear, traced with `perf` to AVX2 gather stalls in the extend kernel; the gather kernels are removed (HPCBio/wfa2lib-rs#31). +- **2026-10-15:** BiWFA overlap scan vectorized (HPCBio/wfa2lib-rs#31): the + breakpoint checks were 73% of affine2p BiWFA time and scalar except the + NEON M scan. affine2p ultralow is now 0.38–0.52× C on ARM (was ~1.0×); + output unchanged. diff --git a/src/bialign.rs b/src/bialign.rs index 68db983..8ac9845 100644 --- a/src/bialign.rs +++ b/src/bialign.rs @@ -111,97 +111,36 @@ pub fn breakpoint_m2m( let ptr_0 = unsafe { wf_0.offsets_centered_ptr() }; let ptr_1 = unsafe { wf_1.offsets_centered_ptr() }; - let mut k_0 = max_lo; - - #[cfg(target_arch = "aarch64")] - { - use std::arch::aarch64::*; - let count = min_hi - max_lo + 1; - if count >= 4 { - unsafe { - let v_tlen = vdupq_n_s32(text_length); - let neon_end = max_lo + (count & !3); - // k_1 = (tlen - plen) - k_0, decreasing: for batch starting at k_0, - // k_1 values are [k_1_0, k_1_0-1, k_1_0-2, k_1_0-3]. - // In memory those are at ptr_1[k_1_0-3..k_1_0], so load and reverse. - let mut k_1_top = wavefront_k_inverse(k_0, pattern_length, text_length); - - while k_0 < neon_end { - let v_off0 = vld1q_s32(ptr_0.offset(k_0 as isize)); - // Load [k_1_top-3, k_1_top-2, k_1_top-1, k_1_top] then reverse to [k_1_top, k_1_top-1, k_1_top-2, k_1_top-3] - let v_raw1 = vld1q_s32(ptr_1.offset((k_1_top - 3) as isize)); - let v_r64 = vrev64q_s32(v_raw1); - let v_off1 = vcombine_s32(vget_high_s32(v_r64), vget_low_s32(v_r64)); - - let v_sum = vaddq_s32(v_off0, v_off1); - let v_mask = vcgeq_s32(v_sum, v_tlen); - - if vmaxvq_u32(v_mask) != 0 { - // Scalar scan within this batch to find first hit - let sums: [i32; 4] = std::mem::transmute(v_sum); - for (i, &sum_val) in sums.iter().enumerate() { - if sum_val >= text_length { - let k_hit = k_0 + i as i32; - let k_1_hit = k_1_top - i as i32; - let offset_0 = *ptr_0.offset(k_hit as isize); - let offset_1 = *ptr_1.offset(k_1_hit as isize); - if breakpoint_forward { - breakpoint.score_forward = score_0; - breakpoint.score_reverse = score_1; - breakpoint.k_forward = k_hit; - breakpoint.k_reverse = k_1_hit; - breakpoint.offset_forward = offset_0; - breakpoint.offset_reverse = offset_1; - } else { - breakpoint.score_forward = score_1; - breakpoint.score_reverse = score_0; - breakpoint.k_forward = k_1_hit; - breakpoint.k_reverse = k_hit; - breakpoint.offset_forward = offset_1; - breakpoint.offset_reverse = offset_0; - } - breakpoint.score = score_0 + score_1; - breakpoint.component = ComponentType::M; - return; - } - } - } - - k_0 += 4; - k_1_top -= 4; + // SAFETY: k_0 ∈ [max_lo, min_hi] ⊆ [lo_0, hi_0] and k_1 ∈ [wf_1.lo, wf_1.hi]. + unsafe { + scan_overlap( + ptr_0, + ptr_1, + max_lo, + min_hi, + pattern_length, + text_length, + |k_0, k_1, offset_0, offset_1| { + if breakpoint_forward { + breakpoint.score_forward = score_0; + breakpoint.score_reverse = score_1; + breakpoint.k_forward = k_0; + breakpoint.k_reverse = k_1; + breakpoint.offset_forward = offset_0; + breakpoint.offset_reverse = offset_1; + } else { + breakpoint.score_forward = score_1; + breakpoint.score_reverse = score_0; + breakpoint.k_forward = k_1; + breakpoint.k_reverse = k_0; + breakpoint.offset_forward = offset_1; + breakpoint.offset_reverse = offset_0; } - } - } - } - - // Scalar tail (or full scalar on non-aarch64) - while k_0 <= min_hi { - let k_1 = wavefront_k_inverse(k_0, pattern_length, text_length); - // SAFETY: bounds guaranteed by construction above. - let offset_0 = unsafe { *ptr_0.offset(k_0 as isize) }; - let offset_1 = unsafe { *ptr_1.offset(k_1 as isize) }; - - if offset_0 + offset_1 >= text_length { - if breakpoint_forward { - breakpoint.score_forward = score_0; - breakpoint.score_reverse = score_1; - breakpoint.k_forward = k_0; - breakpoint.k_reverse = k_1; - breakpoint.offset_forward = offset_0; - breakpoint.offset_reverse = offset_1; - } else { - breakpoint.score_forward = score_1; - breakpoint.score_reverse = score_0; - breakpoint.k_forward = k_1; - breakpoint.k_reverse = k_0; - breakpoint.offset_forward = offset_1; - breakpoint.offset_reverse = offset_0; - } - breakpoint.score = score_0 + score_1; - breakpoint.component = ComponentType::M; - return; - } - k_0 += 1; + breakpoint.score = score_0 + score_1; + breakpoint.component = ComponentType::M; + true + }, + ); } } @@ -244,42 +183,161 @@ pub fn breakpoint_indel2indel( let ptr_0 = unsafe { wf_0.offsets_centered_ptr() }; let ptr_1 = unsafe { wf_1.offsets_centered_ptr() }; - for k_0 in max_lo..=min_hi { + unsafe { + scan_overlap( + ptr_0, + ptr_1, + max_lo, + min_hi, + pattern_length, + text_length, + |k_0, k_1, offset_0, offset_1| { + // Extra bounds check for indel breakpoints + let (check_k, check_h) = if breakpoint_forward { + (k_0, offset_0) + } else { + (k_1, offset_1) + }; + let v = wavefront_v(check_k, check_h); + if v > pattern_length || check_h > text_length { + return false; + } + + if breakpoint_forward { + breakpoint.score_forward = score_0; + breakpoint.score_reverse = score_1; + breakpoint.k_forward = k_0; + breakpoint.k_reverse = k_1; + breakpoint.offset_forward = offset_0; + breakpoint.offset_reverse = offset_1; + } else { + breakpoint.score_forward = score_1; + breakpoint.score_reverse = score_0; + breakpoint.k_forward = k_1; + breakpoint.k_reverse = k_0; + breakpoint.offset_forward = offset_1; + breakpoint.offset_reverse = offset_0; + } + breakpoint.score = score_0 + score_1 - gap_opening; + breakpoint.component = component; + true + }, + ); + } +} + +/// Diagonals checked per vector step in [`scan_overlap`]. +const SCAN_CHUNK: i32 = 8; + +/// Scan the overlap `k_0 ∈ [lo, hi]` in ascending order for diagonals where +/// `ptr_0[k_0] + ptr_1[k_1] >= text_length` (`k_1` is `k_0` in the other +/// direction's coordinates), calling `accept` on each until it returns true. +/// +/// Chunks of [`SCAN_CHUNK`] diagonals are first tested together with SIMD; +/// only a chunk with a candidate is re-checked one diagonal at a time. The +/// visiting order is unchanged, so the breakpoint chosen is the same as a +/// plain scalar scan (and as C's scalar and AVX2 versions). +/// +/// # Safety +/// `ptr_0[k]` for `k ∈ [lo, hi]` and `ptr_1[k_inverse(k)]` must be readable. +#[inline(always)] +unsafe fn scan_overlap( + ptr_0: *const WfOffset, + ptr_1: *const WfOffset, + lo: i32, + hi: i32, + pattern_length: i32, + text_length: i32, + mut accept: impl FnMut(i32, i32, WfOffset, WfOffset) -> bool, +) -> bool { + let mut k_0 = lo; + #[cfg(target_arch = "x86_64")] + let use_avx2 = cfg!(target_feature = "avx2") || std::arch::is_x86_feature_detected!("avx2"); + while hi - k_0 + 1 >= SCAN_CHUNK { + // The chunk's k_1 values run down from k_1_top: they sit at + // ptr_1[k_1_top - 7 ..= k_1_top] in reverse order. + let k_1_top = wavefront_k_inverse(k_0, pattern_length, text_length); + let p_0 = unsafe { ptr_0.offset(k_0 as isize) }; + let p_1 = unsafe { ptr_1.offset((k_1_top - (SCAN_CHUNK - 1)) as isize) }; + #[cfg(target_arch = "aarch64")] + let any = unsafe { chunk_any_neon(p_0, p_1, text_length) }; + #[cfg(target_arch = "x86_64")] + let any = if use_avx2 { + unsafe { chunk_any_avx2(p_0, p_1, text_length) } + } else { + unsafe { chunk_any_scalar(p_0, p_1, text_length) } + }; + #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))] + let any = unsafe { chunk_any_scalar(p_0, p_1, text_length) }; + if any { + for i in 0..SCAN_CHUNK { + let offset_0 = unsafe { *ptr_0.offset((k_0 + i) as isize) }; + let offset_1 = unsafe { *ptr_1.offset((k_1_top - i) as isize) }; + if offset_0 + offset_1 >= text_length + && accept(k_0 + i, k_1_top - i, offset_0, offset_1) + { + return true; + } + } + } + k_0 += SCAN_CHUNK; + } + while k_0 <= hi { let k_1 = wavefront_k_inverse(k_0, pattern_length, text_length); let offset_0 = unsafe { *ptr_0.offset(k_0 as isize) }; let offset_1 = unsafe { *ptr_1.offset(k_1 as isize) }; + if offset_0 + offset_1 >= text_length && accept(k_0, k_1, offset_0, offset_1) { + return true; + } + k_0 += 1; + } + false +} - if offset_0 + offset_1 >= text_length { - // Extra bounds check for indel breakpoints - let (check_k, check_h) = if breakpoint_forward { - (k_0, offset_0) - } else { - (k_1, offset_1) - }; - let v = wavefront_v(check_k, check_h); - if v > pattern_length || check_h > text_length { - continue; - } +/// Whether any `p_0[i] + p_1[7 - i] >= text_length` for `i ∈ 0..8`. +#[cfg_attr(target_arch = "aarch64", allow(dead_code))] +#[inline(always)] +unsafe fn chunk_any_scalar(p_0: *const WfOffset, p_1: *const WfOffset, text_length: i32) -> bool { + let mut any = false; + for i in 0..SCAN_CHUNK as usize { + any |= unsafe { *p_0.add(i) + *p_1.add(SCAN_CHUNK as usize - 1 - i) } >= text_length; + } + any +} - if breakpoint_forward { - breakpoint.score_forward = score_0; - breakpoint.score_reverse = score_1; - breakpoint.k_forward = k_0; - breakpoint.k_reverse = k_1; - breakpoint.offset_forward = offset_0; - breakpoint.offset_reverse = offset_1; - } else { - breakpoint.score_forward = score_1; - breakpoint.score_reverse = score_0; - breakpoint.k_forward = k_1; - breakpoint.k_reverse = k_0; - breakpoint.offset_forward = offset_1; - breakpoint.offset_reverse = offset_0; - } - breakpoint.score = score_0 + score_1 - gap_opening; - breakpoint.component = component; - return; - } +#[cfg(target_arch = "aarch64")] +#[inline(always)] +unsafe fn chunk_any_neon(p_0: *const WfOffset, p_1: *const WfOffset, text_length: i32) -> bool { + use std::arch::aarch64::*; + unsafe { + let v_tlen = vdupq_n_s32(text_length); + // Reverse each 4-lane half of p_1; the high half pairs with the low half of p_0. + let rev = |v: int32x4_t| { + let r = vrev64q_s32(v); + vcombine_s32(vget_high_s32(r), vget_low_s32(r)) + }; + let lo_1 = rev(vld1q_s32(p_1.add(4))); + let hi_1 = rev(vld1q_s32(p_1)); + let m_lo = vcgeq_s32(vaddq_s32(vld1q_s32(p_0), lo_1), v_tlen); + let m_hi = vcgeq_s32(vaddq_s32(vld1q_s32(p_0.add(4)), hi_1), v_tlen); + vmaxvq_u32(vorrq_u32(m_lo, m_hi)) != 0 + } +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +#[inline] +unsafe fn chunk_any_avx2(p_0: *const WfOffset, p_1: *const WfOffset, text_length: i32) -> bool { + use std::arch::x86_64::*; + unsafe { + let rev = _mm256_set_epi32(0, 1, 2, 3, 4, 5, 6, 7); + let v_0 = _mm256_loadu_si256(p_0 as *const __m256i); + let v_1 = _mm256_permutevar8x32_epi32(_mm256_loadu_si256(p_1 as *const __m256i), rev); + let mask = _mm256_cmpgt_epi32( + _mm256_add_epi32(v_0, v_1), + _mm256_set1_epi32(text_length - 1), + ); + _mm256_movemask_epi8(mask) != 0 } } @@ -605,4 +663,111 @@ mod tests { assert_eq!(overlap_gopen_adjust(6, 24, true), 6); assert_eq!(overlap_gopen_adjust(6, 3, true), 3); } + + /// Naive reference for `scan_overlap`: every hit in ascending order. + #[allow(clippy::too_many_arguments)] + fn naive_hits( + a_0: &[i32], + a_1: &[i32], + c_0: i32, + c_1: i32, + lo: i32, + hi: i32, + plen: i32, + tlen: i32, + ) -> Vec { + (lo..=hi) + .filter(|&k| { + let k_1 = wavefront_k_inverse(k, plen, tlen); + a_0[(c_0 + k) as usize] + a_1[(c_1 + k_1) as usize] >= tlen + }) + .collect() + } + + #[test] + fn test_scan_overlap_matches_scalar() { + use crate::offset::OFFSET_NULL; + // Deterministic pseudo-random offsets, some null, around tlen/2. + let mut seed = 0x2545_f491u32; + let mut next = || { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + seed + }; + let (plen, tlen) = (90, 100); + let n = 80usize; + let (c_0, c_1) = (40i32, 40i32); + for trial in 0..400 { + let mut offsets = || -> Vec { + (0..n) + .map(|_| { + let r = next(); + if r % 7 == 0 { + OFFSET_NULL + } else { + 30 + (r % 25) as i32 + (trial % 3) * 3 + } + }) + .collect() + }; + let (a_0, a_1) = (offsets(), offsets()); + // Overlap lengths 1..=26 cover pure tails, whole chunks and both. + let lo = -12 + (next() % 8) as i32; + let hi = lo + (next() % 26) as i32; + let hits = naive_hits(&a_0, &a_1, c_0, c_1, lo, hi, plen, tlen); + // Reject the first `skip` hits, so a rejected hit inside a chunk + // must not end or skip the rest of that chunk. + for skip in 0..3 { + let mut seen = Vec::new(); + let found = unsafe { + scan_overlap( + a_0.as_ptr().offset(c_0 as isize), + a_1.as_ptr().offset(c_1 as isize), + lo, + hi, + plen, + tlen, + |k_0, k_1, o_0, o_1| { + assert_eq!(k_1, wavefront_k_inverse(k_0, plen, tlen)); + assert_eq!(o_0, a_0[(c_0 + k_0) as usize]); + assert_eq!(o_1, a_1[(c_1 + k_1) as usize]); + seen.push(k_0); + seen.len() > skip + }, + ) + }; + let want: Vec = hits.iter().copied().take(skip + 1).collect(); + assert_eq!(seen, want, "trial {trial} lo {lo} hi {hi} skip {skip}"); + assert_eq!(found, hits.len() > skip); + } + } + } + + #[test] + fn test_chunk_any_simd_matches_scalar() { + let a_0 = [10, 20, 30, 40, 50, 60, 70, 80]; + let tlen = 100; + for lane in 0..8 { + for delta in [-1, 0, 1] { + // Only pair (lane, 7 - lane) can reach tlen. + let mut a_1 = [0i32; 8]; + a_1[7 - lane] = tlen - a_0[lane] + delta; + let want = unsafe { chunk_any_scalar(a_0.as_ptr(), a_1.as_ptr(), tlen) }; + assert_eq!(want, delta >= 0); + #[cfg(target_arch = "aarch64")] + assert_eq!( + unsafe { chunk_any_neon(a_0.as_ptr(), a_1.as_ptr(), tlen) }, + want + ); + #[cfg(target_arch = "x86_64")] + if std::arch::is_x86_feature_detected!("avx2") { + assert_eq!( + unsafe { chunk_any_avx2(a_0.as_ptr(), a_1.as_ptr(), tlen) }, + want + ); + } + } + } + } }