Parallelise kinetic electron preconditioner solve using static condensation

[johnomotani]

Work in progress - parallelises preconditioner solve by:

1. using a Schur-complement decomposition to separate the part of the preconditioner action that applies to the electron distribution function from the part that applies to the electron pressure
2. using a static-condensation technique to decompose the distribution-function grid into decoupled blocks that can be solved in parallel, and some 'joining elements' that still require a serial solve, but for a much smaller matrix than the full preconditioner matrix.

Todo:

• Add documentation of the algorithms.
• Improve the description in this PR.
• Register the https://github.com/johnomotani/MPISchurComplements.jl and https://github.com/johnomotani/MPIStaticCondensation.jl packages to simplify installation.
• Further optimization:
  • Extend the Schur-complement decoupling by treating the parallel heat flux to decouple integral-over-g_e terms from the electron distribution function solve, allowing static condensation over both z and w_parallel dimensions. At the moment can only decompose z because the integral terms couple every w_parallel grid point.
  • Apply static condensation technique recursively. Larger blocks mean fewer 'joining elements' so a smaller matrix for the serial solve, so if we first split into few blocks, then split each of those blocks in a new static-condensation solve, etc. we should reduce the serial bottleneck significantly.
  • Apply some straightforward parallelisation of matrix-multiplication in the Schur-complement solve (Parallelise the solve where possible johnomotani/MPISchurComplements.jl#2).

[johnomotani]

Current status with this branch at fbacfc8, johnomotani/MPISchurComplements.jl@4559277, johnomotani/MPIStaticCondensation.jl@90456cc:

• benchmark run examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-coarse_tails-KeCaARK324-teg-wall-Jac-evenvpaelement.toml
  • restarted from the final state of examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-coarse_tails-evenvpaelement-init.toml, followed by examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_boltzmann-coarse_tails-evenvpaelement.toml.
  • compare to serial run with LU solver, using examples/kinetic-electrons/wall-bc_recyclefraction0.5_split3_kinetic-coarse_tails-KeCaARK324-teg-wall-Jac-evenvpaelement-lu.toml
• On 32 core AMD Threadripper workstation, serial LU run takes 498 s. This new solver takes: 1 core, 718 s; 2 cores, 419 s; 4 cores, 295 s; 8 cores, 240 s; 16 cores 195 s; 32 cores 204 s.
  • There are only 32 elements in z in this run, so on 32 cores there is only one z-element per block in the static condensation grid decomposition, which is a challenging case!

Installation

At the moment, until the MPISchurComplements.jl and MPIStaticCondensation.jl packages are registered, you will need to clone those repos and dev them in the top-level moment_kinetics repository before switching to this branch (or before starting to set up moment_kinetics). In a fresh clone of moment_kinetics, this would look something like (updating the paths to the other repos depending on where you put them):

moment_kinetics$ touch Project.toml
moment_kinetics$ julia --project -O3
julia>]
(moment_kinetics) pkg> dev ../MPISchurComplements.jl
(moment_kinetics) pkg> dev ../MPIStaticCondensation.jl
(moment_kinetics) pkg> ^D
moment_kinetics$ ./machines/machine_setup.sh

and follow through the prompts as usual to set up moment_kinetics.

[johnomotani]

It is expected that all the CI tests fail at the moment, because the MPISchurComplements.jl and MPIStaticCondensation.jl packages are not registered, and in anticipation that they will be eventually we have not put special setup in the Github Actions config files to install them.