speed up the 128-bit ksw2_extd2/extz2 inner loops - #64
Conversation
|
I will have a look later. Does the change keep the identical output? |
|
Yes — identical, verified at the kernel and at the aligner. Kernel, per-pair exact. Both kernels are compiled twice into one binary (baseline from
That gate is not vacuous: flipping a single direction bit in the candidate ( Aligner, byte-identical SAM. Rebased onto
Rep-to-rep spread is under 0.3% in every cell (e.g. stock WGS wall 56.99/57.08/57.11 s), and the ranges are non-overlapping throughout. Two robustness checks on those numbers. The same digests come out of a second machine (M2 Max) at 8 threads rather than 16, so the output is invariant to host and thread count. And the ratios are unchanged at 1/10th the input size — single-copy runs gave 1.037/1.101/1.081× against 1.038/1.110/1.086× here — so the concatenation is not flattering the result through a warm cache. Kernel throughput (M3 Ultra, arm64/NEON, n=20000, 7 reps), 128-bit baseline → candidate:
One caveat worth stating: the short-read gain is the least reliable figure here. It is small, and unlike the long-read numbers it did not reproduce across hosts — the M2 Max measured 1.07× where this machine measures 1.04×. Treat WGS as "a few percent, host-dependent"; the HiFi and ONT figures agreed across both machines to within 0.3%. 🤖 Measurements run and this comment drafted with Claude Code (Claude Opus 5); reviewed before posting. |
|
Thanks. Will have a look when I have time |
Four changes to the shared anti-diagonal loop, all output-neutral. The per-diagonal score prepass compared against three constants to pick the match, mismatch and N scores. It now indexes a 16-byte table with sf ^ qrr and one byte shuffle, 5 ops down to 2. Query N is remapped to 8 so every index stays <= 12, where NEON's vqtbl1q_u8 and SSSE3's _mm_shuffle_epi8 agree. Each DP rail was shifted one lane with srli+slli+or. Carrying the whole previous vector instead of its last byte makes that a single vext (alignr on x86), three rails in the dual-gap kernel and two in the single-piece one. The exact-max path widened four int8 to int32 at a time; it now does sixteen, keeping the four-lane accumulator and the update order so max_t is unchanged. The direction byte was built with four and/orr pairs and a blend. NEON spells each as one BIT, which clang will not emit from the portable form: vbslq_u8(k, g, d) lowers to (k & g) | (~k & d), and for a constant splat known-bits proves ~k & d == d and folds it back. The instruction is written out for arm64, with the plain or/and and blendv kept elsewhere -- the flag sites need d's k bits already clear, the argmax site is a real blend, and sharing one form between them would be wrong off arm64. Together the left-alignment core loop goes from 77 to 67 instructions per 16 cells on aarch64. Measured on a 16-core Graviton4 at -t 16, mapping phase only: HiFi 10k 8.425s -> 7.480s, ONT 10k 9.909s -> 8.892s, WGS 1M pairs 9.907s -> 9.624s, Hi-C 1M pairs 73.766s -> 73.076s. Long reads gain most because bw_long leaves their gap fills effectively unbanded, so extd2 is most of their runtime. SAM output is byte-identical to r421 on 200k WGS pairs and 10k HiFi reads, and a differential harness comparing the old and new kernels directly reports 60000 pairs identical in both score and CIGAR at qlen 100, 150 and 250.
bc0384a to
a917828
Compare
2 participants
Four changes to the shared anti-diagonal loop. Output is unchanged.
The per-diagonal score prepass compared against three constants to pick the match,
mismatch and N scores. It now indexes a 16-byte table with sf ^ qrr and one byte
shuffle, 5 ops down to 2. Query N is remapped to 8 so every index stays <= 12,
where NEON's vqtbl1q_u8 and SSSE3's _mm_shuffle_epi8 agree.
Each DP rail was shifted one lane with srli+slli+or. Carrying the whole previous
vector instead of its last byte makes that one vext on NEON and one palignr on
SSSE3 -- three rails in extd2, two in extz2.
The exact-max path widened four int8 to int32 at a time; it now does sixteen,
keeping the four-lane accumulator and the update order so max_t is unchanged.
The direction byte was built with four and/orr pairs and a blend. NEON spells each
as one BIT, which clang will not emit from the portable form: vbslq_u8(k, g, d)
lowers to (k & g) | (~k & d), and for a constant splat known-bits proves
~k & d == d and folds it back. The instruction is written out for aarch64;
elsewhere the original or/and and blendv are kept. Two wrappers rather than one,
because the flag sites need d's k bits already clear while the argmax site is a
real blend, and sharing one form between them would be wrong off aarch64.
The first three help any SIMD target; the fourth is aarch64-only. Left-alignment
core loop, 77 to 67 instructions per 16 cells on aarch64.
Timings use the four HG002 sets from zenodo.org/records/19703025, the same ones
map.mak drives, aligned against GRCh38 with ALT contigs and decoys rather than
hs38. Mapping phase only, -t 16, SAM to /dev/null, five reps per arm interleaved
with order reversal:
HiFi 10k 8.425s 7.480s +11.2 11.316s 10.539s +6.9
ONT 10k 9.909s 8.892s +10.3 14.078s 13.148s +6.6
WGS 1M pairs 9.907s 9.624s +2.9 14.976s 14.515s +3.1
Hi-C 1M pairs 73.766s 73.076s +0.9 - - -
Rep spread was 0.01-0.77%. Long reads gain most because bw_long leaves their gap
fills effectively unbanded, so extd2 is most of their runtime; on short reads it is
a much smaller share. The x86 column is a baseline -msse4.2 build, so it does not
include whatever the AVX2/AVX-512 dispatch in #20 would add.
SAM output is byte-identical to r421 on 200k WGS pairs and 10k HiFi reads, on both
architectures. A differential harness that builds the old and new kernels into one
binary and compares them per pair reports 60000 pairs identical in score and CIGAR
at qlen 100, 150 and 250, for both extd2 and extz2.