Enhanced preconditioner for JOREK MHD solver

JOREK [1, 2] is a state-of-the-art code for the simulation of large-scale instabilities in tokamaks (magnetic confinement fusion devices). It solves extended reduced [3] or full [4] MHD equations using a 2D iso-parametric G¹ continuous finite element discretization [5], which was recently extended to arbitrary order [6], combined with a toroidal Fourier decomposition. Main applications are simulations addressing the dynamics of major disruptions including vertical displacement events (VDEs) and runaway electrons as well as edge localized modes (ELMs) [2]. The goal is to advance the physics understanding for these potentially very harmful phenomena and to develop reliable control schemes applicable in present and future fusion devices.

A free boundary and resistive wall extension exists [7] and a framework for kinetic effects is available [8]. The time integration is fully implicit such that a very large—and, due to the properties of the equations, badly conditioned—sparse matrix system needs to be solved in every time step. The present article describes improvements of the computational efficiency of the code by a variety of different approaches enhancing the solver performance. These methods may also be beneficial in other codes.

1.1.1. JOREK solver.

Starting with the general form of MHD equations

$$\begin{align} \frac{\partial\mathbf{A}(\mathbf{u})}{\partial t} = \mathbf{B}(\mathbf{u},t), \end{align}$$

we get the following linearized with respect to $\delta\mathbf{u}$ time discretization

$$\begin{align} & \left[\left(1+\zeta\right)\left(\frac{\partial\mathbf{A}}{\partial\mathbf{u}}\right)^n - \Delta t\theta\left(\frac{\partial\mathbf{B}}{\partial\mathbf{u}}\right)^n\right]\delta\mathbf{u}^n \nonumber\\ &\quad = \Delta t\mathbf{B}^n + \zeta\left(\frac{\partial\mathbf{A}}{\partial\mathbf{u}}\right)^{n-1}\delta\mathbf{u}^{n-1}. \end{align} \tag{ 1 }$$

Here, n represents the time step index, and $\delta\mathbf{u} = \delta\mathbf{u}^{n+1}-\delta\mathbf{u}^n$ where u is a vector of all unknowns. Setting $\theta = 1/2$ and ζ = 0, equation (1) recovers the Crank–Nicolson scheme, while alternatively setting θ = 1 and $\zeta = 1/2$ the BDF2 (Gears) scheme is recovered.

After spatial discretization, equation (1) forms a large linear algebraic system of equations

$$\begin{align} A\,x = b \end{align} \tag{ 2 }$$

where A is a sparse matrix. Two main options are available for that: when the problem size is small, in particular for axisymmetric cases, a direct sparse matrix solver is applied to the complete system. The second option is applied for the large problem sizes typically encountered in real-life applications: the system is solved iteratively with the generalized minimal residual (GMRES) method [9].

Due to the stiffness of the system, e.g. caused by the highly anisotropic diffusion operator and the presence of large convective terms in MHD equations the matrix A is ill-conditioned, i.e. a small change in the right hand side b leads to a big change of the solution x. This causes extremely poor convergence of the GMRES algorithm. To improve the convergence, (left-) preconditioning (PC) is applied [10], namely the original system (2) is replaced by

$$\begin{align} M^{-1}A\,x = M^{-1}b \end{align} \tag{ 3 }$$

where M is a non-singular preconditioner matrix, which should be close to A in some sense, and the solution

$$\begin{align} z = M^{-1}w \end{align} \tag{ 4 }$$

should be easy to find in comparison to solving the original system. At the first step of the GMRES iteration cycle w = b, and later on $w = b-Ax_m$, where x_m is the solution of (3) at mth iteration. It is noteworthy that the matrix M is only involved in the algorithm for obtaining the solution vectors.

In the standard JOREK algorithm, the preconditioner matrix M has block-diagonal structure with each block corresponding to the individual toroidal Fourier modes. Physically this corresponds to the approximation of A by omitting mode-coupling terms. A schematic representation for a small case with the toroidal modes $n = (0,1,2,3)$ is shown in figure 1. The solution vector z is then formed by combining the individual solutions for each diagonal block, which can be found independently by using direct sparse solver libraries. The advantage of this method is its scalability with respect to the toroidal resolution and robustness: the underlying physics leading to ill-conditioning is addressed by using the direct solver applied to the PC matrices.

Zoom In Zoom Out Reset image size

The PC matrix M can usually be reused for many time steps thus saving on the expensive LU-factorization of the blocks. In the linear phase of the physical system's evolution, when the mode amplitudes are small and mode coupling is negligible, the preconditioner matrix M is an almost exact approximation of the full matrix A such that the iterative solver usually converges in a single iteration (we typically chose the relative backward error to be smaller than 10⁻⁷ as the convergence criteria). However, when the problem becomes non-linear such that the coupling between toroidal modes cannot be neglected any more, the preconditioner becomes far less efficient, and consequently the GMRES convergence deteriorates. This poor approximation of A by the PC matrix M is one of the main challenges for the efficiency of the JOREK solver. Besides that, the large memory consumption of the LU decomposed preconditioner matrix blocks and the limited parallel scalability of the direct sparse matrix solvers are additional challenges to be faced.

In the present work, we address these challenges, leading to a significant step forward in code efficiency. This includes a complex representation of preconditioner block-matrices (section 2.1) and a generalization of the preconditioner allowing to approximate the matrix A more accurately in highly non-linear scenarios (section 2.2). Properties of corresponding matrices are described in the appendix. The performance of the new developments is evaluated in section 3.

Now that the JOREK solver has been introduced, we identify the major numerical stages performed during a single time step. The stages are abbreviated by one letter in the following to simplify notation in figures and tables later on. Note, that the present article concentrates on the phases after the global matrix construction. Here, F and I are usually most critical for the overall performance.

• G: Global matrix construction. Matrix A is constructed in a distributed way at every time step. This can cost a significant fraction of the overall run time depending on the case.
• P: Construction of PC block-matrices. So far done by extracting from the global matrix via communication; an alternative option for direct construction was implemented. This part berries a significant computational cost comparable to G, however it is not usually performed at every time step thus not as critical.
• A: Analysis of the PC matrix structure. Here matrix is analyzed and elimination tree is build. This is only done in the very first time step of a simulation, since the structure is not changing during a simulation, but only the values of the matrix entries. This phase is therefore usually not critical for the overall performance.
• F: LU-factorization of the PC matrices. Done only in time steps where the preconditioner needs to be updated. Often critical for the overall performance. The relative costs compared to phase I depend on the application and time step used, but also on the user defined threshold for updating the PC matrix—the user can influence the frequency of the preconditioner updates to some degree via the choice of the time step and via a threshold.
• I: Iterative solution of the global system including solver calls for the preconditioner matrix blocks. Done in every time step. Often critical for the overall performance. The relative costs compared to the preconditioner updates of phase F depend on case and settings (see explanation above).

The quantities listed in the following are used in this article to describe the benchmark setups and results.

• N_N: Total number of compute nodes used in the simulation
• N_M: Total number of MPI tasks used in the simulation
• M_T: Total memory consumption for the whole simulation (sum for all MPI tasks)
• t_W: Elapsed wall clock time
• $t_\textrm{N}\equiv t_\textrm{W}\cdot N_\textrm{N}$: Computational cost in node-hours

All benchmarks performed in this work are carried out on the Marconi supercomputer operated by CINECA. The partition Marconi–Fusion available to the European fusion research community comprises of 2396 compute nodes, each equipped with two 24-core Intel Xeon 8160 CPUs ('Skylake' generation) at a base clock speed of 2.10 GHz and 192 GB DDR4 RAM. The interlink between the compute nodes is provided by Intel OmniPath (100 Gbit s⁻¹). We use Intel v.2018.4 compilers with compatible intelmpi and mkl libraries.

The rest of the article is organized as follows. Section 2 describes the developments performed as part of the present work. Section 2.1 explains the adaptation to complex PC matrices and section 2.2 describes a generalization of the preconditioner, which improves convergence in highly non-linear simulation cases. In section 3, previously existing and new solver options are compared based on realistic simulation cases to investigate performance improvements obtained in the present work. Finally, section 4 briefly provides summary, conclusions, and an outlook to future developments.

In order to optimize performance of a direct solver we have exploited the idea of representing the preconditioner block-matrices in complex form. This reduces the memory consumption considerably and improve overall solve time. The transformation into complex format is only possible when the matrix has the appropriate symmetry structure. For example, the following set of equations in the matrix form

$$\begin{align} \begin{pmatrix} a & -b \\ b & a \end{pmatrix} \cdot \begin{pmatrix} x \\ y \end{pmatrix} = \begin{pmatrix} c \\ d \end{pmatrix} \end{align} \tag{ 5 }$$

can be replaced by

$$\begin{align} \begin{pmatrix} a + ib \end{pmatrix} \cdot \begin{pmatrix} x + iy \end{pmatrix} = \begin{pmatrix} c + id \end{pmatrix}. \end{align} \tag{ 6 }$$

Since numerical symmetry is not exactly preserved, the following symmetrization is performed:

$$\begin{align} \begin{pmatrix} a & -b \\ c & d \end{pmatrix} \rightarrow \begin{pmatrix} \left(a+d\right)/2 & -\left(b+c\right)/2 \\ \left(b+c\right)/2 & \left(a+d\right)/2 \end{pmatrix}. \end{align} \tag{ 7 }$$

This symmetrized matrix can now be replaced by a complex form, analogously to (6). As will be shown later, the approximation introduced by this symmetrization affects the efficiency of the preconditioner only very mildly (small impact on the number of GMRES iterations) such that the overall benefit of using complex matrices prevails. It is to note here that the symmetry enforcement is a crucial step and without it, a convergence can not be achieved. Besides switching to complex input, nothing else was changed in the direct solver library interface. Performance benchmarks of the complex representation of PC matrices are presented in section 3.1.

In this section, we describe our approach for improving the efficiency of the preconditioner in highly non-linear scenarios, where the coupling between the toroidal Fourier modes becomes strong causing the GMRES convergence to deteriorate massively or break down entirely. The approach followed here, is to combine several toroidal modes into 'mode groups'. The interaction between the modes within the mode groups is retained in the PC matrix M. Consequently, larger block-matrices need to be solved than in the original preconditioner, which makes the solution more expensive in terms of memory consumption and computing time. Nevertheless, capturing the non-linear interaction can improve the convergence of the iterative solver so strongly that the overall performance is significantly enhanced, as demonstrated in section 3.2.

While the proposed approach is extremely flexible in terms of choosing mode groups, we illustrate it using the following example of overlapping mode groups (see figure 2). In this example, referred to as Overlapping, the preconditioner block-matrices consist of entries from the global matrix A corresponding to a pair of toroidal Fourier modes. Each mode is thus present in two consecutive mode groups. Additional single-mode blocks are added corresponding the first and the last modes. When the preconditioner solution is formed, each row is filled with the linear combination of the solutions from two corresponding blocks. the weighting factors can be chosen freely. Note, that we do not assemble a complete preconditioner matrix M explicitly, just a preconditioner solution (4) as needed for the GMRES iterations. In the described approach we effectively take into account interactions of neighboring modes which physically resemble to mode coupling as observed in turbulence cascading, thus we denote this preconditioner as physics based.

Zoom In Zoom Out Reset image size

Another alternative approach we have considered in this work couples the typically strongest n = 0 mode with all other $n\ne0$ harmonics by forming $\{(0),\,(0,1),\,(0,2),\dots\}$ mode groups.

The implementation is performed in a very general way allowing more advanced configurations for mode groups than the examples given. It is, for instance, possible to combine modes separated by a certain toroidal period, modes which are believed to have the strongest interactions. Also an automatic and dynamic selection of the groups could be realized in the future. The implementation in JOREK is done as follows:

• Via the namelist input file, arbitrary mode groups can be constructed in a flexible way. Each mode group forms a diagonal block in the preconditioner matrix.
• The number of MPI tasks for each mode group can be automatically assigned according to the number of non-zero entries in each mode group, or manually be chosen via the input file. This allows optimizing the load balance for the non-uniform matrix block sizes encountered here.
• Mode groups can be overlapping. When reconstructing the whole preconditioner solution $z = M^{-1}w$ as linear combination of the individual solutions from the mode group blocks, factors are applied. For example, the factor is $1/2$ in the Overlapping case.

Mode group methods are benchmarked in section 3.2 with the standard preconditioner.

The benchmark is based on a simplified simulation of ELMs in a JET-like geometry, i.e. an X-point plasma with open and closed field lines included in the simulation domain. A fairly high poloidal grid resolution of 51k Bezier finite elements is used, while only the toroidal mode numbers $n = (0,4,8)$ are included in the simulation. Tests for this case are performed both in the linear phase, where the exponentially growing instabilities are too small to modify the axisymmetric n = 0 component and in the non-linear phase during the collapse of the edge pressure gradient caused by the instability.

Table 1 contains data from the simulations restarted in the highly nonlinear regime for a single time-step on six and three compute nodes with two MPI tasks per node. For the given problem size, the real PaStiX [11, 12] solver needs at least six compute nodes based on the memory requirement. The complex PaStiX solver, on the other hand, can be used on three compute nodes due to the reduced memory consumption.

In table 1, we compare the wall-clock time t_W (sec) spent on solver phases using six compute nodes for the real and six respectively three compute nodes for the complex PaStiX solver.

Due to the reduced matrix dimension, factorization (F) as one of the most expensive parts of the solve step is roughly 2.5 and 3.3 times faster when using the complex solver on six and three compute nodes, respectively, compared to the real solver. The analysis (A) is also done more efficiently with the complex solver, however, being performed only in the first time step of a simulation, it adds limited benefits for a typical application. Moreover, one can notice from the table 1 that although the complex solver requires one additional GMRES iteration as compared to the real solver, the GMRES solve (I) on six respectively three compute nodes is more efficient by a factor of 1.2 and 1.8 as compared to the real solver on six compute nodes.

Figure 3 shows the time-lines of the memory consumption for the real solver (with six compute nodes) and complex solver (with six and three compute nodes). It is apparent that the maximum memory consumption when using the complex solver with six nodes is roughly 600 GB, which is less than roughly 700 GB in the case of the real solver. The memory consumption is even lower in the case on complex solver using three compute nodes. Thus, we confirm that the complex solver reduces the overall memory consumption significantly. It is worth mentioning that we observe a larger reduction in memory consumption for a higher number of toroidal modes.

Zoom In Zoom Out Reset image size

In conclusion, it can be seen from table 1, that the overall performance gain by the complex solver is moderate when the memory consumption imposes the same minimum number of compute nodes for the real and complex solvers and when updates of the PC matrix (factorizations) are needed seldomly. On the other hand, if the PC matrix is updated more often, the complex solver may reduce the total computational cost (including matrix construction and preconditioner matrix assembly which are not explicitly discussed in this section) by a third. When the reduced memory consumption by the complex solver allows to use a lower number of compute nodes, the overall computational costs can be approximately halved by moving into a regime where the scalability of the solver improves.

The simulation considered here is a slightly adapted version of the simulations presented in [13]. It represents a VDE, that develops violent non-axisymmetric ($n\ne0$) instabilities while moving into the plasma facing components. We have chosen this case, since the standard iterative solver convergence is poor, forcing the use of extremely small time steps. It is suspected that the strong non-linearity of the problem reduces the accuracy by which the PC matrix M approximates the complete system A, such that condition number and convergence deteriorate.

We have used this VDE case to test the new preconditioner based on the mode groups. For this purpose we run simulations in the nonlinear regime, and compare the Standard preconditioner (figure 1) with our new approach based on Overlapping mode groups (figure 2). Another mode-group based approach couples the n = 0 dominant harmonic with all other modes, referred to as n = 0 coupling.

As shown in table 2, the Overlapping preconditioner reduces the number of GMRES iterations by a factor of 3, leading to the reduction of the whole GMRES cycle time by a factor of 2.2, respectively. Meanwhile, the factorization time increased by a factor of 6. Thus, the overall performance depends on how often preconditioner updates are required. A similar improvement trend with the Overlapping method is observed for the VDE case with higher toroidal resolution including $n = (0,1,2,\dots,10)$, as shown in table 3.

The analysis, factorization and solve times with the new preconditioner are determined by the corresponding times spent on the largest diagonal block, which is twice the size of the standard case. Since each diagonal block-matrix is factorized and solved independently, the difference in matrix sizes introduces a strong workload imbalance. To address this, we have implemented a flexible distribution of MPI tasks to each mode group. For example, in the current simulation using the Overlapping method, we have five mode groups and a total of 32 MPI tasks. The groups consist of the following toroidal modes: $\{(0),(0,1),(1,2),(2,3),(3)\}$. The block size corresponding to n = 0 mode is twice smaller than for any other mode, and the number of nonzero entries is four times smaller. The preconditioner block-matrix sizes for each group are shown in table 4.

There are different ways to distribute the available 32 MPI tasks among the mode groups. For instance, distributing them as (2,6,9,9,6) would increase the factorization time to 147.6 s as compared to 123.9 s in the case of (1,4,12,12,3) tasks per group distribution. Here one should take into account that the most computationally expensive factorization has a different MPI task scaling compared to the GMRES iteration cycle, thus the optimal task distribution for balanced performance depends on the relative contribution of total factorization time and total iteration time. However, as a rule of thumb, one can distribute tasks roughly proportionally to the number of nonzero entries in each mode group. In $n = (0,1,2,3)$ we have assigned 1 and 3 MPI tasks to the single mode groups of the (0) and the (3) modes, respectively, four tasks to the $(0,1)$ group, and 12 tasks to each of the $(1,2)$ and the $(2,3)$ groups. In the $n = (0,1,2,\dots,10)$ case 2 MPI tasks are assigned to the smallest mode group of (0) mode, six MPI tasks are assigned to the mode group of (11) mode, and the overlapping mode groups get eight MPI tasks each.

Since the PC matrix needs not to be updated (factorized) in every time step, the efficiency of the new preconditioner is determined by the actual frequency of these updates. In our example, the Overlapping preconditioner would lead to the same overall performance, if the preconditioner is updated every five time steps. This number would be slightly higher in practice, since the time to construct preconditioner matrices is also slightly higher for the new preconditioner. In the nonlinear phase of the VDE simulation with a normalized time step of 0.05, the preconditioner matrix updates need to be performed approximately every 50 time steps, which gives the new preconditioner a strong advantage.

Meanwhile, as the preconditioner approximates the total system a lot better now, it becomes possible to increase the simulation time step allowing for an even more efficient way of running the simulation. For the same nonlinear case, good overall performance was obtained when increasing the time step by a factor of 3 for the Overlapping preconditioner, leading to the total t_W spent on simulating 100 steps using eight compute nodes of 104 min. Comparing this to 97 min required to perform 100 steps with $\Delta t = 0.05$ using the Standard preconditioner, an approximate speed-up of a factor 3 × is obtained. For the n = 0 coupling preconditioner the time step could be increased by a factor of 2 leading to the total simulation time of 95 min, i.e. a factor 2 × speed-up compared to the Standard approach. The same conclusion holds for the VDE case with a higher toroidal resolution including the $n = (0,\dots,10)$ toroidal modes, which was tested as well.

The main disadvantage of the new preconditioner lies in the increased memory consumption. As we can see in the figures 4 and 5, for the VDE case with $n = (0,1,2,3)$, the maximum required memory with the Overlapping preconditioner is ∼75% higher than for the Standard method, while for the n = 0 coupling preconditioner the overall memory consumption increases by ∼30%.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

The presented results demonstrate just an example of using the new preconditioner. With the provided implementation flexibility, there is room for further performance improvements, e.g. by adjusting work load balance, data locality, optimization of GMRES preconditioner update frequency vs. simulation time step, etc.