GPU acceleration

[cjfields]

Explore whether key steps can be optionally run using GPU if one is present, in particular NW alignment steps used across learn-errors, dada/dada-pooled, and remove-bimera-denovo.

[cjfields]

Scoped plan

Discussion narrowed the GPU target to NW alignment in learn-errors and dada_uniques specifically, since both funnel through the same hot loop.

Single kernel covers both paths

Both paths go through:

• b_compare_parallel (src/cluster.rs:118) — called per cluster birth in dada_uniques; aligns every Raw against the new cluster center
• build_trans_mat (src/learn_errors.rs, called per learn-errors iteration) — re-aligns each input against its assigned cluster center to count transitions

Both have the "one center vs N queries" shape — the center sequence is fixed across the batch. Ideal GPU pattern: stage the center to shared memory once, parallelize over queries. One kernel implementation accelerates both call sites, and therefore learn-errors, dada, and dada-pooled simultaneously.

Order-of-magnitude scale (F3D0 Illumina):

• ~131 clusters × ~2000 raws ≈ 260K alignments per dada_uniques call
• learn-errors runs ~5 iterations → ~1.3M alignments per learn-errors run, plus build_trans_mat re-alignments

Three design constraints to plan around

1. K-mer screen rejects >90% of pairs today (`raw_align_with_buf`, src/nwalign.rs:1047). Either run the screen on GPU too (better occupancy, simpler dataflow — likely the right call) or screen host-side and ship only survivors (preserves existing filter logic but pays PCIe per call and shrinks batches).

2. Greedy short-circuit at src/cluster.rs:124 (`if greedy && (raw.reads > center_reads || raw.lock)`) skips the alignment entirely. Data-dependent branching wastes GPU lanes. Pre-partition queries host-side into "skip" vs "align" sets.

3. `compute_lambda` runs alongside the alignment in the same parallel pass. Fuse lambda into the GPU kernel rather than shipping `Sub` structs back across PCIe. `err_mat` is small (16 × nq ≈ 6 KB) and stable across an entire dada run — upload once per run.

Deployment

Behind an optional Cargo feature flag:

```toml
[features]
default = []
cuda = ["dep:cudarc"]
```

• Default build is CPU-only, no CUDA in the dependency tree
• `--features cuda` links cudarc, runtime auto-detects GPU, falls back to CPU if absent
• `--no-gpu` CLI flag for users who want to force CPU on CUDA-built binaries

Sanity check before committing to this

Instrument `b_compare_parallel` to log nraw, n_greedy_skipped, n_kmer_rejected per call. If median "n_aligned" per call is:

• < ~1000 → kernel launch + PCIe overhead likely eat the speedup at A40 scale
• > ~10000 → sweet spot, GPU is worth it
• In between → depends on alignment length (PacBio HiFi at ~1500 bp tolerates smaller batches than 240 bp Illumina)

PacBio HiFi typically has fewer raws than Illumina (longer reads → fewer uniques), so measuring on the actual target workload matters more than F3D0 numbers.

Prerequisite

#18 calls for a flamegraph on real PacBio HiFi to confirm alignment is still the dominant cost after the over-budding fix landed in #4. That measurement gates the GPU work — if alignment is <20% of runtime on PacBio, this isn't worth the engineering investment.

Out of scope (for now)

• `remove-bimera-denovo` — alignment-heavy but called less frequently than dada. Worth revisiting only after the dada/learn-errors path is GPU-accelerated.
• K-mer screen as a standalone GPU step — too fast on CPU to justify PCIe overhead.
• Loess fitting, dereplication, taxonomy — wrong shape for GPU.

[cjfields]

Reassessment after the NW kernel vectorization (still worth it, but gate on workload scale)

Following the alignment-kernel work (see #18: dploop vectorization, diag-fill range-hoist, u8 pointers — the kernel is now ~1.6× faster than R DADA2 per alignment, and learn-errors on a small PacBio sample went 18.6s → 7.8s). Revisiting whether GPU is still worth pursuing:

Alignment is still the dominant cost — there's large remaining headroom

The vectorization captured a big chunk but didn't change the shape: alignment was ~93% of runtime, and after a ~2.4× kernel speedup it's still ~85% (Amdahl: 93/2.4 + 7 ≈ 46 total, alignment 39/46). So the theoretical ceiling from making alignment free is still ~6.7× more end-to-end. GPU-accelerating the NW step is not a marginal gain.

A point that now favors GPU more, not less

The CPU kernel is memory-bandwidth-bound — the d/p DP matrices streaming through memory (confirmed: u8 pointers only bought ~2%). On GPU this inverts: with band=32 the active band is ~65 cells, small enough to keep per-alignment state in registers/shared memory, sidestepping the exact global-memory bottleneck the CPU can't escape. So a batched banded-NW GPU kernel could be compute-bound and scale better than a naive 10× estimate. (Caveat unchanged: traceback needs the full p matrix → still global memory or recompute-in-traceback.)

What's changed against GPU

• CPU baseline rose 2.4×, so GPU's relative win is smaller than vs the old code.
• rayon scales well (6.8× on 8 cores) — "add CPU cores" is a real cluster alternative.
• Engineering cost unchanged and significant: CUDA banded-NW kernel with DADA2's exact scoring (+ homopolymer variant), the gapless screen, lambda fusion, cuda feature flag, CPU fallback. Weeks, not days.

The deciding factor: real-workload scale

The earlier cluster runs were ~11.5h (grain-sweep table) — far larger than the 7.8s small-sample bench, so the real samples are much bigger. Re-measure a representative full-size run with the current build first.

• 11.5h → ~5h from the kernel win alone, and still painful → GPU justified (alignment is 85%, GPU removes the memory bottleneck; plausibly tens of minutes).
• New kernel brings it to acceptable → GPU not worth the engineering cost.

Recommendation

Don't start GPU yet — gate it on one measurement. Re-run a full-size sample (or the real multi-sample job) on the cluster with the post-#18 build and see where the multi-hour runs land. That number decides it. The Amdahl headroom (~6.7×) plus the GPU-sidesteps-memory-bound argument mean GPU is genuinely worth it if the scale still hurts — but the CPU vectorization may have already closed enough of the gap to make it unnecessary.

Scoping for if/when we proceed remains in the earlier comment (kernel covers b_compare_parallel + build_trans_mat, fused lambda, behind a cuda feature flag).

[cjfields]

Benchmark target: use a problematic (poorly-dereplicating) sample, not a typical one

Refining the "gate on a real-workload measurement" point above.

Alignment cost scales roughly as O(clusters × uniques × iterations). Samples that dereplicate inefficiently (many uniques surviving — low-complexity, high-error, or contaminated libraries) hit a multiplicative blowup: more uniques to align against each cluster center, often more clusters, compounded across self-consistency iterations — and that inflated alignment cost then slows every downstream step (build_trans, shuffle, etc.).

Implications:

• A typical/clean sample would understate the GPU case (alignment finishes quickly, overhead doesn't dominate).
• The decisive benchmark is an actual problematic sample — the worst-case regime is exactly where a fixed CPU kernel (even now that it beats R) stops keeping up and a throughput device pays off.

So when we gate the GPU decision on a measurement, the target should be a known-problematic sample at full scale, capturing: total alignments, wall-clock, and alignment's fraction of runtime. If alignment overhead skyrockets there even with the post-#18 kernel, that's the signal GPU is worth the engineering cost; if even the hard samples are tolerable, it isn't.

TODO when revisited: identify a representative poorly-dereplicating PacBio sample and run the post-#18 build on it (full scale, real thread count) before committing to GPU work.