Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use Kokkos-Kernels ODE solvers instead of custom ones #851

Open
BenWibking opened this issue Jan 14, 2025 · 1 comment
Open

use Kokkos-Kernels ODE solvers instead of custom ones #851

BenWibking opened this issue Jan 14, 2025 · 1 comment
Labels
enhancement New feature or request radiation

Comments

@BenWibking
Copy link
Collaborator

BenWibking commented Jan 14, 2025

Describe the proposal
It should clean up the code a lot if we switch to using the Kokkos-Kernels ODE solvers (https://github.com/kokkos/kokkos-kernels/tree/develop/ode) and Newton-Raphson solver (https://github.com/kokkos/kokkos-kernels/blob/develop/ode/impl/KokkosODE_Newton_impl.hpp).

They are well-designed and cleanly separate the "physics" (the right-hand-side) from the "math" (the ODE solver itself). The advantage of these solvers is that are all templated on the number of equations to solve (plus other parameters, like the floating point datatype), so we don't need to set any compile-time parameters like we have to do for Microphysics.

For example, the Newton solver can be called for each cell like this:

    KokkosODE::Experimental::newton_solver_status newton_status =
        KokkosODE::Experimental::Newton::Solve(sys, param, jac, tmp_gesv, y_new, rhs, update, scale);

Big caveat: this would add a dependency on Kokkos, which might be tricky, since Kokkos is designed to interact with the GPU in a similar way to AMReX. We need to check whether it's possible to call Kokkos::ODE from within amrex::ParallelFor.

Alternatively, we could AMReXify the Kokkos-Kernels source code: we need to replace the KOKKOS_* macros with their AMREX_* equivalents, and replace Kokkos:: functions with AMReX equivalents, and then copy these headers into src/math. The difficult part is that we also need to do this for the Kokkos BLAS functions, which are highly nontrivial.

Describe alternatives you've considered
Keep using our custom ODE solvers.

Additional context
We can use the Kokkos ODE solvers to replace:

  • the Runge-Kutta ODE solver for cooling
  • the Newton-Raphson solver we use for radiation, and
  • the VODE ODE solver we use for chemistry
@BenWibking BenWibking added enhancement New feature or request radiation labels Jan 14, 2025
@BenWibking
Copy link
Collaborator Author

@chongchonghe This might be of interest.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request radiation
Projects
None yet
Development

No branches or pull requests

1 participant