[Experimental / Engine]: Add optional CUDA FDTD execution backends - #192
Draft
QuentinWach wants to merge 114 commits into
Draft
[Experimental / Engine]: Add optional CUDA FDTD execution backends#192QuentinWach wants to merge 114 commits into
QuentinWach wants to merge 114 commits into
Conversation
QuentinWach
changed the title
Freeze each CPML timestep behind a second XLA-owned field bank and alternate source-aware H/E launches between the two banks. This removes same-buffer field hazards and opens the regular-grid core to safe H-to-E fusion. Advance beamz-cuda to ABI v6 / package 0.6.0 for the new typed FFI target; retain an environment switch for paired diagnostics. RTX 3090 modal waveguide, CPML-12, 106x172x172 (3.14M cells), 192 steps: two alternating 9-sample pairs measured 7.279/7.273 GCUPS versus 6.514/6.514 GCUPS for the legacy in-place schedule (+11.7%). Active peak memory increases from 0.51 to 0.58 GiB. Validation: 44 ABI/runtime unit tests and 26 CUDA hardware parity tests passed (2 Hopper-only skips); Ruff and diff checks pass.
Submit the disjoint absorbing shell and plain deep-core work through one mixed CUDA block queue. This lets GA102 fill branch and memory latency across the two workloads while removing one graph node from each H/E phase; retain an environment switch for exact paired controls. RTX 3090 modal waveguide, CPML-12, 106x172x172 (3.14M cells), 192 steps: two alternating 11-sample pairs measured 7.463/7.458 GCUPS versus 7.393/7.393 GCUPS for serial shell/core launches (+0.91%). All 26 applicable CUDA hardware parity tests pass (2 skips); 44 ABI/runtime unit tests, Ruff, and diff checks pass.
Keep alternating field banks across arbitrary source and packed DFT monitor programs, eliminating the per-step field preservation traffic while retaining exact FP32 CPML state semantics. Add hardware parity coverage and an RTX 3090 benchmark profile for CPML with two sources and two monitors. RTX 3090, 96x128x192, 200 steps, CPML-8, two sources, two plane DFT monitors: Temporal graph: 6.587 / 6.409 / 6.403 GCUPS. In-place graph: 6.110 / 5.977 / 5.985 GCUPS. Paired gains: +7.81% / +7.23% / +6.98%. Validation: 45 unit tests and CPML multi-monitor CUDA/JAX hardware parity pass.
Route eligible source-free CPML simulations through the existing temporal source-group target with empty groups. This adds no graph kernels and avoids duplicating a native handler while removing in-place field preservation traffic. RTX 3090, 96x128x192, 200 steps, heterogeneous CPML-8: temporal 7.816 / 7.865 GCUPS versus old 7.202 / 7.179 GCUPS, paired gains +8.53% / +9.56%. Validation: 24 CUDA runtime contract tests and even/odd complete-state CUDA/JAX hardware parity pass.
Freeze both field and FP32 CPML recurrence inputs across each temporal graph step by allocating a second compact psi bank and alternating it with the field banks. Bump the native ABI to 8 for the expanded temporal output contract. RTX 3090, 106x172x172 (3,135,904 cells), 192 steps, 12-cell CPML, modal source, 11 samples / 4 warmups: - temporal psi: 7.921 and 7.869 GCUPS - in-place psi control: 7.242 GCUPS - improvement: +9.4% and +8.7% Validation: 46 CUDA runtime/backend unit tests; 27 hardware parity tests passed (2 skipped).
Run exact FP32 CPML field updates and arbitrary source timings inside one cooperative CUDA grid with grid-wide leapfrog barriers. Keep field and psi ping-pong banks frozen, avoid per-step graph dispatch, and preserve coincident and overlapping source semantics. Use the persistent scheduler for monitor-free regular CPML programs from 16 steps through 10 million cells; retain the graph path above the measured RTX 3090 crossover or when cooperative launch is unavailable. RTX 3090 canonical modal waveguide, 192 steps, 12-cell CPML, eight source terms, 11 samples / 4 warmups: - 106x172x172 (3,135,904 cells): 11.026 GCUPS - graph control: 8.127 GCUPS - improvement: +35.7% - 182x296x296 (15,946,112 cells): gated graph path, 8.762 GCUPS Crossover sweep retained 9.405 GCUPS at 9,625,000 cells; the first losing point was 10,816,000 cells. Validation: 46 CUDA runtime/backend unit tests and 27 hardware parity tests passed (2 skipped).
Limit cooperative CPML launches to two billion cell-steps as well as ten million cells. Longer jobs retain the CUDA-graph scheduler so an interactive RTX display is not monopolized by one non-preemptible multi-step grid. The 3.14M-cell, 192-step 11.026-GCUPS workload remains eligible; the full 27-test hardware parity suite passes (2 skipped).
Keep scheduled, ragged, and windowed DFT accumulation inside the cooperative multi-step kernel, sharing phase/window evaluation across each 32-point warp. Specialize monitor-free launches so their 48-register occupancy and 11.08 GCUPS modal result are preserved.\n\nRTX 3090, 96x128x192, 200 steps, CPML + 2 sources + 2 monitors:\n- persistent monitor path: 7.887 GCUPS (59.807 ms median)\n- graph replay control: 6.854 GCUPS (68.805 ms median)\n- improvement: 15.1%\n\nValidation: 27 CUDA hardware tests passed (2 skipped); 46 CUDA runtime/backend tests passed.
Map resident work at warp granularity so the cooperative schedule covers every x/y tile instead of launching a 1-D block against a 64x4 mapping. Preserve global parity across bounded launch chunks and add a strict full-amplitude resident-vs-graph state test.\n\nThe previous 11 GCUPS resident result was invalid because only one tile row was updated. With full coverage, resident execution measures 4.03 GCUPS on the 3.14M-cell modal CPML case versus 8.13 GCUPS for graph replay. Keep the correct resident scheduler behind BEAMZ_CUDA_ENABLE_PERSISTENT_CPML=1 and restore graph replay as the production default.
Keep the BF16 absorber shell on 32x4 tiles while remapping the same 128 threads to a 64x2 recurrence-free core tile. FP32 retains its measured-optimal 64x4 queue. Add a hardware coverage/parity gate against the independent split shell/core schedule. RTX 3090 modal CPML, 12-cell absorber, 192 steps: the 106x172x172 case improved from about 69.85 ms / 8.62 GCUPS to 69.34 ms / 8.68 GCUPS. At 182x296x296 (15,946,112 cells), an isolated run reached 337.11 ms / 9.082 GCUPS; a thermally loaded repeat was 342.02 ms / 8.952 GCUPS. Validation: 29 hardware tests passed (2 skipped), 46 CUDA runtime/backend unit tests passed.
Carry CUDA scheduling and precision choices through the compiled program and typed FFI instead of consulting mutable process environment state from native launches. Include the policy in program and CUDA graph identities, reject unsupported CUDA sharding early, and consolidate native launch initialization.\n\nValidated on RTX 3090: 81 passed, 2 skipped across focused unit and CUDA hardware tests. The realistic 96x128x192 CPML+source+monitor case reached a 67.67 ms median over 200 steps (6.97 GCUPS), versus 69.29 ms (6.81 GCUPS) before this refactor.
Centralize common FFI attributes, field replacement, source-group packing, and temporal CPML workspace aliases. Replace numeric layout offsets with named constants and remove 63 net lines from the CUDA Python runtime without changing the native execution path.\n\nRTX 3090 realistic CPML+source+monitor benchmark: 67.50 ms median for 200 steps on 96x128x192 (6.99 GCUPS), retaining the preceding refactor's performance.
Decode FFI buffer lists once and construct streamed and temporal CPML launches through shared, named layouts. This removes 207 net lines from the native bridge while keeping XLA ownership and aliasing explicit.\n\nRTX 3090 realistic CPML+source+monitor benchmark: 67.80 ms median for 200 steps on 96x128x192 (6.96 GCUPS). CUDA hardware suite: 29 passed, 2 skipped.
Remove the redundant single-source and single-source-monitor Python, FFI, launch, and kernel stacks. Route every supported source/DFT configuration through the grouped program interface, with internal one-member source and one-monitor specializations and shared native launch helpers. This deletes 738 net lines while preserving arbitrary source timings, ragged/scheduled monitors, PEC, and temporal CPML semantics. Advance beamz-cuda to ABI v10 / package 0.10.0 and refresh the architecture documentation.\n\nPaired RTX 3090, 96x128x192, 200 steps, 15 samples / 4 warmups:\n- CPML + source: 67.897 ms / 6.950 GCUPS -> 58.822 ms / 8.022 GCUPS (+15.4%)\n- CPML + source + DFT: 69.238 ms / 6.815 GCUPS -> 61.086 ms / 7.724 GCUPS (+13.4%)\n- PEC + source: 52.832 ms / 8.931 GCUPS -> 52.397 ms / 9.006 GCUPS (+0.8%)\n- PEC + source + DFT: 54.348 ms / 8.682 GCUPS -> 54.118 ms / 8.719 GCUPS (+0.4%)\n\nRTX 3090 modal waveguide, 182x296x296, 192 steps, CPML-12, BF16 recurrence state:\n- mode source: 342.065 ms / 8.951 GCUPS\n- mode source + four-field, three-frequency DFT monitor: 348.164 ms / 8.794 GCUPS\n\nValidation: 46 CUDA runtime/backend unit tests passed; 29 CUDA hardware tests passed (2 skipped); Ruff, diff, extension manifest, and host FFI syntax checks passed.
The cooperative resident scheduler remained opt-in because it was materially slower than bounded CUDA graph replay on GA102. Remove its cross-cutting source, CPML, monitor, and policy code while retaining the winning temporal graph path.\n\nFocused validation: 46 unit tests and 28 CUDA hardware tests passed (2 skipped). On the RTX 3090 96x128x192/200-step check, CPML+source+DFT measured 7.936 GCUPS versus the 7.690 baseline and PEC measured 13.528 versus 13.325 GCUPS; neither workload regressed.
Decode every multi-step FFI layout into a typed BeamzProgramLaunch, validate phases, sources, monitors, and temporal eligibility once, and centralize CPML geometry construction. Remove obsolete variant launch wrappers and their duplicated validation.\n\nFocused validation: 46 unit tests and 28 CUDA hardware tests passed (2 skipped). RTX 3090 96x128x192/200-step medians remained stable at 7.932 GCUPS for CPML+source+DFT (7.936 before) and 13.557 GCUPS for PEC (13.528 before).
Give graph identity, bounded caching, nested-capture fallback, instantiation, replay, and cleanup one native owner. Build cache keys from explicit program fields instead of raw C++ struct bytes and make every scheduler consume the complete BeamzProgramLaunch.\n\nFocused validation: 46 unit tests and 28 CUDA hardware tests passed (2 skipped). RTX 3090 96x128x192/200-step medians were 7.911 GCUPS for CPML+source+DFT versus the original 7.690 baseline and 13.358 GCUPS for PEC versus 13.325 originally.
Split the former 1,989-line streamed implementation into update/CPML, source/monitor, program scheduling, and graph execution modules. Leaf kernel launchers no longer own graph or timestep policy, while the program layer validates and schedules complete simulations.\n\nRTX 3090, 96x128x192, 200 steps (9 samples):\n- CPML + source + DFT monitor: 60.959 ms median, 7.740 GCUPS (baseline 61.357 ms, 7.690 GCUPS)\n- uniform PEC: 35.243 ms median, 13.388 GCUPS (baseline 35.411 ms, 13.325 GCUPS)\n\nFocused validation: 46 unit tests and 28 hardware tests pass (2 hardware skips).
Make cuda/abi_layout.json the source of truth for ABI versions, target names, and positional buffer layout constants. Generate matching Python and C++ definitions and check them in focused contract tests, removing the duplicated runtime, extension, build, and decoder constants.\n\nRTX 3090, 96x128x192, 200 steps (9 samples):\n- CPML + source + DFT monitor: 59.628 ms median, 7.913 GCUPS\n- uniform PEC: 35.204 ms median, 13.404 GCUPS\nBoth remain ahead of the original 7.690 and 13.325 GCUPS baselines.\n\nValidation: 49 unit tests and 28 hardware tests pass (2 hardware skips).
Replace seven public multi-step target names with one beamz_cuda_program entry and a generated, explicit buffer-layout code. Keep phase and Hopper handlers separate because they have distinct contracts. This cuts the required streamed manifest from eight targets to two and makes layout selection visible in every cached JAX program.\n\nRTX 3090, 96x128x192, 200 steps (9 samples):\n- CPML + source + DFT monitor: 59.478 ms median, 7.933 GCUPS\n- uniform PEC: 35.049 ms median, 13.463 GCUPS\nBoth remain ahead of the original 7.690 and 13.325 GCUPS baselines.\n\nValidation: 49 unit tests and 28 hardware tests pass (2 hardware skips).
Stop copying two complete phase descriptors into each monitor kernel. Pass the six sampled field buffers and timestep directly, reducing Ampere constant-parameter usage for AccumulateDftGroups from 12,172 bytes to 1,236 bytes (89.8%) without changing registers or monitor arithmetic.\n\nRTX 3090, 96x128x192, 200 steps (15 samples):\n- CPML + source + DFT monitor: 61.276 ms median, 7.700 GCUPS\n- uniform PEC: 35.235 ms median, 13.391 GCUPS\nBoth remain ahead of the original 61.357 ms / 7.690 GCUPS and 35.411 ms / 13.325 GCUPS baselines.\n\nHardware parity: 28 passed, 2 skipped.
Give the fused full-step kernel a purpose-built phase payload containing only its 12 Yee inputs, three metrics, three outputs, and inverse resolution. This separates the stable host ABI from the device command and reduces Ampere constant-parameter usage from 11,584 bytes to 2,096 bytes (81.9%) without changing register use.\n\nRTX 3090, 96x128x192, 200 steps (15 samples):\n- uniform PEC: 34.509 ms median, 13.674 GCUPS (13.325 baseline, +2.6%)\n- CPML + source + DFT monitor: 59.272 ms median, 7.961 GCUPS (7.690 baseline, +3.5%)\n\nHardware parity: 28 passed, 2 skipped.
Name the packed boundary attribute accurately, derive remaining decoder capacities and offsets from generated layout constants, remove hard-coded source group counts from native program validation, and document the three-target/module architecture and regeneration workflow.\n\nRTX 3090, 96x128x192, 200 steps (15 samples):\n- CPML + source + DFT monitor: 59.678 ms median, 7.906 GCUPS\n- uniform PEC: 34.400 ms median, 13.716 GCUPS\nBoth remain ahead of the original 7.690 and 13.325 GCUPS baselines.\n\nValidation: generated ABI check, 49 unit tests, and 28 hardware parity tests pass (2 hardware skips).
Install the native module as beamz._cuda, mark its build-only wheel as Private :: Do Not Upload, remove the stale exact BeamZ dependency, and keep CI artifacts private. Validate the generated component version and in-package CMake install contract. The component was rebuilt and imported successfully on the RTX 3090 with ABI v11 and all three FFI registrations.
Generate the two remaining CUDA options and all source/monitor slots from the shared schema. Promote the validated temporal, combined-CPML, source-batching, adaptive-tile, and material-codebook paths to unconditional behavior; remove their legacy switches and dead fallback kernels. Add exact native dtype validation and direct BF16 application parity against JAX. RTX 3090 hardware result: 28 passed, 2 SM90-only skipped.
Bump the portable, RTX comparison, matrix, and capacity schemas so records retain field and CPML precision, resolved backend, BeamZ/JAX/JAXLIB/Python versions, native component version, ABI, and immutable flags. Reject mixed-precision or policy-mismatched comparisons. Real RTX 3090 child smokes emitted beamz._cuda 0.12.0, ABI 12, FP32 CPML, and flags 128.
Preserve the 12-cell default while centralizing the invariant that numerical profile sampling only sees a resolved physical thickness. This closes the Pyright trust-gate errors without changing profile arithmetic.
Keep the CUDA extension intrinsic to BeamZ by including its private build component in the BeamZ sdist only. Strengthen distribution validation so local caches cannot leak and the component remains excluded from published BeamZ wheels and independent package publication.
QuentinWach
force-pushed
the
t3code/research-cuda-fdtd-acceleration
branch
from
August 14, 2026 20:17
d4df8fc to
0d91c1b
Compare
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
1 participant
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.
Summary
This PR adds an optional, CUDA-accelerated execution path for three-dimensional FDTD while preserving BeamZ's public simulation API and JAX backend.
beamz._cudamodule rather than publishing a separate CUDA productNote
The streamed SM80–SM89 backend is validated on an RTX 3090. The explicit
cuda_hopperSM90 path remains experimental until it passes the same parity, sanitizer, and throughput gates on Hopper hardware.How it works
The Python compiler lowers the same immutable BeamZ simulation into either JAX operations or a typed JAX FFI program. The CUDA program combines:
The FFI boundary is ABI v12. A single JSON schema generates the Python indices and C++ constants, and the native implementation is split into decoding, program scheduling, graph execution, update/CPML, source/monitor, and Hopper modules.
RTX 3090 results
Warm, already-compiled execution on the same RTX 3090 and JAX/JAXLIB 0.9.0 environment:
origin/mainJAXFor the directly compared CPML/source/monitor workload, the 95% speedup interval is 1.808–1.885× and executable compilation is 7.18× faster. The 15.95-million-cell modal workload runs at 1.754 ms/step with a 2.79 GiB active GPU-memory peak.
GCUPS is
prod(grid_zyx) × timesteps / median warm runtime / 1e9. Timings include field updates, boundaries, configured sources, and configured monitors, but exclude setup and compilation.Correctness and validation
Packaging and fallback policy
CUDA is an optional native part of BeamZ:
beamzbeamz._cudaPrivate :: Do Not Uploadbackend="auto"selects CUDA only for a compatible single-GPU 3D plan and otherwise keeps the JAX pathCurrent limits