Add GPU (wgpu) Direct Fourier transform and trajectory timing#13
Merged
Conversation
Add an optional `gpu` cargo feature providing a wgpu/WGSL implementation of the Direct Fourier transform (vacuum path) for fast trajectory averaging. The kernel evaluates F(q) = Σ f_i·exp(i q·r_i), |F(q)|² in f32 with one workgroup per q-vector and a shared-memory reduction over atoms; static q-vectors and atom ids stay device-resident across frames. When built with the feature, `direct` runs on the GPU by default and `--cpu` forces the CPU path; it falls back to the CPU on an active solvent model or when no adapter is available. Degenerate sizes and device-side failures return errors rather than panicking. Measured ~15x over the rayon CPU path on an Apple M4 (88,400 atoms, 1300 q-vectors), matching the CPU f64 reference to <5e-6 relative. Also add offset-scan and per-frame timing logs to trajectory averaging.
There was a problem hiding this comment.
Pull request overview
Adds an optional GPU-accelerated (“wgpu”/WGSL) implementation of the vacuum Direct Fourier transform to speed up trajectory averaging, along with CLI/library plumbing and additional timing instrumentation.
Changes:
- Introduces
IntensityScheme::DirectGpuwith feature-gated GPU backend and CPU fallback logic. - Adds
--cpuCLI flag to force CPU direct transform when GPU support is compiled in. - Adds GPU implementation (
src/gpu.rs) plus trajectory offset-scan/averaging timing logs.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
src/trajectory.rs |
Adds timing logs around XTC offset scanning and per-frame averaging loop. |
src/lib.rs |
Adds DirectGpu scheme and dispatch logic with feature gating and CPU fallbacks. |
src/gpu.rs |
New wgpu/WGSL direct structure-factor implementation and a GPU-vs-CPU golden test. |
src/explicit.rs |
Exposes average_duplicates for reuse by the GPU path. |
src/cli.rs |
Adds --cpu and maps CLI “direct” to DirectGpu by default when available. |
Cargo.toml |
Adds optional GPU dependencies and a gpu feature flag. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+278
to
+279
| // Per-frame host work: form-factor table, packed positions, box. | ||
| let ff_table = self.form_factor_table(&qmags)?; |
Comment on lines
+207
to
+211
| let elapsed = average_start.elapsed().as_secs_f64(); | ||
| log::info!( | ||
| "Averaged {num_frames} frames in {elapsed:.2} s ({:.1} ms/frame)", | ||
| 1e3 * elapsed / num_frames.max(1) as f64 | ||
| ); |
added 2 commits
June 18, 2026 14:36
Time trajectory averaging against the number of frames actually consumed rather than the offset-table length, which overstates frames (and skews ms/frame) if XtcIterator stops early on a frame read error. Document why the GPU q-magnitudes stay f64 (form factors and q-grid match the CPU path; only the in-shader phase is f32).
Note the optional `gpu` build feature: GPU-accelerated `direct` scheme for trajectory averaging, automatic CPU fallback, and the `--cpu` opt-out.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add an optional
gpucargo feature providing a wgpu/WGSL implementation of the Direct Fourier transform (vacuum path) for fast trajectory averaging. The kernel evaluates F(q) = Σ f_i·exp(i q·r_i), |F(q)|² in f32 with one workgroup per q-vector and a shared-memory reduction over atoms; static q-vectors and atom ids stay device-resident across frames.When built with the feature,
directruns on the GPU by default and--cpuforces the CPU path; it falls back to the CPU on an active solvent model or when no adapter is available. Degenerate sizes and device-side failures return errors rather than panicking.Measured ~15x over the rayon CPU path on Apple M4 (88,400 atoms, 1300 q-vectors), matching the CPU f64 reference to <5e-6 relative.