When the JFNK code in nonlinear_solvers was originally written, loops that could be parallelised were parallelised using the existing shared-memory parallel-loop infrastructure, to simplify the initial implementation using already-tested code. This choice means that the parallelisation depends on the number and kind ($r$, $z$, $w_\parallel$ etc.) of the dimension of the array(s) being solved for - this requires multiple versions of the distributed_norm(), distributed_dot(), parallel_map(), and parallel_delta_x_calc() functions.
As (most of*) the parallel operations in the JFNK solver do not need to know about the structure of the solution - just treating the solution array as a flattened 1D vector - the nonlinear_solvers module could be simplified by instead defining its own parallelisation scheme. Distributed-memory parallelism would stay the same as the current version. Shared-memory parallelism, instead of using the setup used elsewhere (https://mabarnes.github.io/moment_kinetics/dev/developing/#Parallelization), could use a simpler scheme: given a 'shared memory communicator' containing nshared processes, divide the part of the solution vector owned by each block/subdomain into nshared pieces, giving each process in the shared memory communicator a unique range of 'locally owned' indices in the vector, which that process will loop over in every parallelised loop. This locally-owned range would be stored in the nl_solver_info object created by setup_nonlinear_solve().
* Some of the variants in nonlinear_solvers are to deal with cases where we are solving for two variables at once (a 1D1V/1D2V shape function and a 1D pressure). In these cases the 'norm' used by the solver is the RMS of the norms of each of the two variables individually (e.g. here) - it is more important that every entry of the pressure is accurately solved than that every entry of the shape function (some of which are very small values far down in the tails) is. I have not tested what the effect would be of switching this to a single norm on the full solution vector, but I think it seems like a sensible thing to do, so it would be nice to preserve the feature if possible. At present, the solution vector in these cases is stored as a Tuple of arrays, rather than a single array - maybe it makes sense to generalise this approach so that every solution vector is a Tuple of arrays (when a single array is passed it could be wrapped in a one-element Tuple), and initialise/store a corresponding Tuple of 'locally-owned ranges'? This has the bonus that for multiple-variable cases, the arrays containing the variables can be passed directly, without needing to copy them into a combined buffer array.
Notes on a few things I happen to have thought of already:
- We probably can and need to remove the
outer_coords argument to setup_nonlinear_solve() and the precon_lowerz_vcut_inds/precon_upperz_vcut_inds fields in the nl_solver_info struct. outer_coords is only used to initialise those fields, and the fields should probably live somewhere else, like in a preconditioner struct - @johnomotani can look into where they should go.
- There may well be other fields in
nl_solver_info that should be moved outside the nonlinear_solvers module because they relate to how the JFNK solver is used elsewhere in the code, rather than being necessary to the solver itself.
- The
serial_solve, anysv_region, and anyzv_region arguments to setup_nonlinear_solve() should be replaced by just passing in 2 MPI communicators called something like distributed_comm and shared_comm, where distributed_comm only needs to be defined on rank-0 of shared_comm.
- The
electron_p_pdf_solve argument to setup_nonlinear_solve() can be replaced by having generic multiple-variable support as discussed in (*) above.
- The
coords argument to setup_nonlinear_solve() should probably be replaced by an integer or Tuple-of-integers argument giving the size(s) of the solution vector (or solution vector components when there are multiple variables). We can work these out in moment_kinetics in the places where setup_nonlinear_solve() is called.
When the JFNK code in$r$ , $z$ , $w_\parallel$ etc.) of the dimension of the array(s) being solved for - this requires multiple versions of the
nonlinear_solverswas originally written, loops that could be parallelised were parallelised using the existing shared-memory parallel-loop infrastructure, to simplify the initial implementation using already-tested code. This choice means that the parallelisation depends on the number and kind (distributed_norm(),distributed_dot(),parallel_map(), andparallel_delta_x_calc()functions.As (most of*) the parallel operations in the JFNK solver do not need to know about the structure of the solution - just treating the solution array as a flattened 1D vector - the
nonlinear_solversmodule could be simplified by instead defining its own parallelisation scheme. Distributed-memory parallelism would stay the same as the current version. Shared-memory parallelism, instead of using the setup used elsewhere (https://mabarnes.github.io/moment_kinetics/dev/developing/#Parallelization), could use a simpler scheme: given a 'shared memory communicator' containingnsharedprocesses, divide the part of the solution vector owned by each block/subdomain intonsharedpieces, giving each process in the shared memory communicator a unique range of 'locally owned' indices in the vector, which that process will loop over in every parallelised loop. This locally-owned range would be stored in thenl_solver_infoobject created bysetup_nonlinear_solve().* Some of the variants in
nonlinear_solversare to deal with cases where we are solving for two variables at once (a 1D1V/1D2V shape function and a 1D pressure). In these cases the 'norm' used by the solver is the RMS of the norms of each of the two variables individually (e.g. here) - it is more important that every entry of the pressure is accurately solved than that every entry of the shape function (some of which are very small values far down in the tails) is. I have not tested what the effect would be of switching this to a single norm on the full solution vector, but I think it seems like a sensible thing to do, so it would be nice to preserve the feature if possible. At present, the solution vector in these cases is stored as a Tuple of arrays, rather than a single array - maybe it makes sense to generalise this approach so that every solution vector is a Tuple of arrays (when a single array is passed it could be wrapped in a one-element Tuple), and initialise/store a corresponding Tuple of 'locally-owned ranges'? This has the bonus that for multiple-variable cases, the arrays containing the variables can be passed directly, without needing to copy them into a combined buffer array.Notes on a few things I happen to have thought of already:
outer_coordsargument tosetup_nonlinear_solve()and theprecon_lowerz_vcut_inds/precon_upperz_vcut_indsfields in thenl_solver_infostruct.outer_coordsis only used to initialise those fields, and the fields should probably live somewhere else, like in a preconditioner struct - @johnomotani can look into where they should go.nl_solver_infothat should be moved outside thenonlinear_solversmodule because they relate to how the JFNK solver is used elsewhere in the code, rather than being necessary to the solver itself.serial_solve,anysv_region, andanyzv_regionarguments tosetup_nonlinear_solve()should be replaced by just passing in 2 MPI communicators called something likedistributed_commandshared_comm, wheredistributed_commonly needs to be defined on rank-0 ofshared_comm.electron_p_pdf_solveargument tosetup_nonlinear_solve()can be replaced by having generic multiple-variable support as discussed in (*) above.coordsargument tosetup_nonlinear_solve()should probably be replaced by an integer or Tuple-of-integers argument giving the size(s) of the solution vector (or solution vector components when there are multiple variables). We can work these out inmoment_kineticsin the places wheresetup_nonlinear_solve()is called.