Skip to content

feat(vlasov1d): add native CUDA/MPI backend with kinetic ions - #389

Merged
joglekara merged 9 commits into
mainfrom
codex/native-vlasov-ions
Oct 27, 2026
Merged

joglekara merged 9 commits into
mainfrom
codex/native-vlasov-ions

Conversation

@joglekara

@joglekara joglekara commented Oct 26, 2026 •

Copy link
Copy Markdown
Member

Add backend: native for the existing Vlasov-1D ergoExo lifecycle and custom BaseVlasov1D subclasses. C++/CUDA/MPI advances independent kinetic species with spectral x streaming, ADEPT's cubic C1 velocity interpolation, sixth-order splitting, a shared Poisson field, and optional reference-species Chang-Cooper Dougherty collisions. Existing Python initialization, JAX save callbacks, xarray/NetCDF postprocessing and MLflow artifacts are retained; the solve does not require AD.

Native workers expose synchronized DLPack views and stream saved outputs into host memory maps. Standard saves use actual timesteps, with configurable scalar cadence. Unsupported physics and incompatible shared libraries fail explicitly. Include the CMake build, MPI/FFTW CI coverage, configuration reference, and historical A100 comparison artifacts.

Validation: 8 local native/observation tests and 211 current-main lifecycle/core/builder/configuration regressions pass. Coverage includes 1/2-rank MPI, coupled electron/ion/alpha grids and charge/mass ratios, custom callbacks, ergoExo logging, and NetCDF roundtrips. The built wheel includes all native build sources.

Five A100 integration cases passed: 100-step kinetic electron/proton cases on 1/4/8 GPUs (max saved-value errors 1.12e-12, 9.31e-13, 9.66e-13), a 4-GPU electron regression, and the unchanged current turbulence main module on 8 GPUs including f0, velocity-spectrum and cascade callbacks (max error 1.01e-11). Its postprocessor wrote all seven expected NetCDFs, reopened successfully. Per-leaf comparison uses rtol 1e-10 / atol 1e-11 at actual save times. The GPU allocation is released.

The Oct26 validation report is included; its raw source archives, NetCDF files and MLflow metadata remain local/on NERSC. Collisions retain ADEPT's reference-species convention; this does not implement electron-ion collisional exchange. Native execution uses ergoExo; the explicit JAX builder rejects native config to prevent silent fallback.

Historical warmed physics timings against pinned ADEPT/JAX measured 1.28x on 4 A100s and 2.74x on 8 A100s. These exclude preprocessing, compilation and output and are not new ion performance claims.

Review corrections use resolved envelope defaults for both constant-profile validation and collision frequency, and expose the allocation's selected CUDA ordinal separately from the MPI local rank (ABI3). Native synchronization restores the allocation's device. Omitted time/space pulse defaults and constant-baseline defaults are covered. The reviewer approved both corrections; both inline threads are resolved.

Post-review validation: 11 focused CPU tests pass locally and in test-native-vlasov1d for implementation commit 1441e51. A 100-step proton case on two distinct GPUs with one GPU visible per rank passes (max error 8.31e-13); current turbulence callbacks also pass in that mode (max error 5.74e-12). Both MPI local ranks 0/1 correctly use CUDA ordinal zero. On this shared Cray node, the masked regression uses host MPI and socket NCCL with P2P/SHM disabled; the all-visible production launcher is unchanged. The final e767f67 commit changes only the harness, documentation and written report; solver/test sources are identical to 1441e51. All 15 full checks passed before these isolated native review fixes; repeated broad suites may still be running.

@joglekara
joglekara marked this pull request as ready for review October 26, 2026 21:55

@ergodic-reviewer ergodic-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Requesting changes for two correctness defects: omitted collision-envelope defaults can silently change the configured physics, and MPI local ranks are incorrectly used as CUDA device ordinals when each worker sees only one GPU. Details are attached to the affected code.

Postprocessing is implemented and reachable: native saves retain the existing callbacks and actual time axes, feed the multispecies field/scalar/distribution plotting and NetCDF paths, and pass through ergoExo artifact logging. The added tests exercise multispecies NetCDF roundtrips and custom-callback artifacts; no duplicate postprocessor is needed.

Static review of dd58cf9bec6214f49cba99ebeea90d32877f54cc. I did not execute tests, benchmarks, or postprocessing. The published GPU report was inspected; its external raw plots/NetCDF artifacts were not available for independent inspection.

Comment thread adept/_vlasov1d/native/module.py
Comment thread adept/_vlasov1d/native/worker.py Outdated

@ergodic-reviewer ergodic-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The previous blocking findings are addressed: collision validation and the constant rate now use resolved envelope values, with omitted-default regression tests; worker JAX selection and DLPack metadata now use the selected CUDA ordinal rather than the MPI local rank. No additional material findings from static review.

Postprocessing remains implemented and reachable through the existing lifecycle: native callbacks produce multispecies moments, scalar diagnostics and distributions with actual save times, consumed by the plotting/NetCDF writers and ergoExo artifact logging. Tests cover multispecies NetCDF roundtrips and custom-callback artifacts; no duplicate postprocessor is needed.

GitHub reports test-native-vlasov1d successful for this head; other suites were still running when checked. I did not execute tests, benchmarks or postprocessing. The published GPU validation report was inspected, but its external raw plots/NetCDF artifacts were not independently accessible.

@ergodic-reviewer ergodic-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No material findings from static review. The latest change retains the allocation's CUDA device ordinal and restores it before native synchronization, keeping DLPack metadata and subsequent native work tied to the owning device. The earlier collision-default and per-worker GPU-visibility fixes remain in place.

Postprocessing is implemented and reachable: native callbacks retain multispecies moments, scalar diagnostics and distributions with actual save times, consumed by the existing plotting/NetCDF writers and ergoExo artifact logging. Tests cover multispecies NetCDF roundtrips and custom-callback artifacts; no duplicate postprocessor is needed.

Validation: inspected the implementation, tests and published GPU report. I did not execute tests, benchmarks or postprocessing, and the external raw GPU artifacts were not independently accessible. For this head, GitHub reports successful base, installation and pre-commit checks; native/Vlasov tests and several other suites are still running.

@ergodic-reviewer ergodic-reviewer Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No material findings from static review at e767f672fa101672363739b3933c3ae142f59087. The native integration preserves the existing initialization/save/postprocessing boundaries, explicitly rejects unsupported physics and training, and retains the earlier collision-default and CUDA-device fixes.

Postprocessing is implemented and reachable: native save callbacks produce multispecies moments, scalar diagnostics and distributions, consumed by the existing plotting/NetCDF writers and ergoExo artifact logging. Standard saves carry actual timestep axes; custom saves can preserve requested times through state interpolation. Tests cover multispecies NetCDF roundtrips and custom-callback artifacts; no duplicate postprocessor is needed.

Validation was limited to source, test and report inspection. I did not execute tests, benchmarks or postprocessing, and did not independently inspect the external raw GPU plots/NetCDF artifacts.

@joglekara
joglekara merged commit c0a8ef0 into main Oct 27, 2026
16 checks passed
@joglekara
joglekara deleted the codex/native-vlasov-ions branch October 27, 2026 00:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant