Landau Collision Integral Solver with Adaptive Mesh Refinement on Emerging Architectures
M. F. Adams E. Hirvijoki M. G. Knepley J. Brown T. Isaac R. Mills




Inversion of the exciton built-in dipole moment in In(Ga)As quantum dots via nonlinear piezoelectric effect
    Rinaldo Trotta
    
===========================================================================================================




  The Landau collision integral is an accurate model for the
  small-angle dominated Coulomb collisions in fusion plasmas.  We
  investigate a high order accurate, fully conservative, finite
  element discretization of the nonlinear multi-species Landau
  integral with adaptive mesh refinement using the PETSc library
  (www.mcs.anl.gov/petsc).  We develop algorithms and techniques to
  efficiently utilize emerging architectures with an approach that
  minimizes memory usage and movement and is suitable for vector
  processing.  The Landau collision integral is vectorized with Intel
  AVX-512 intrinsics and the solver sustains as much as 22% of the
  theoretical peak flop rate of the Second Generation Intel Xeon Phi,
  Knights Landing, processor.



  Landau collision integral, fusion plasma physics




§ INTRODUCTION

  
The simulation of magnetized plasmas is of commercial and scientific
interest and is integral to the DOE's fusion energy research program
<cit.>.







Although fluid models are widely employed to model fusion plasmas, the
weak collisionality and highly non-Maxwellian velocity distributions in such plasmas motivate
the use of kinetic models, such as the so-called Vlasov-Maxwell-Landau system. The evolution of the phase-space density or
distribution function f of each species (electrons and multiple
species of ions in general) is modeled with

    df/dt≡∂ f/∂ t + ∂ x/∂ t·∇_x f+ ∂ v/∂ t·∇_v f
    = ∂ f/∂ t +  v·∇_x f+
    e/m(  E +  v× B) ·∇_v f = C

where e is charge, m mass, E electric field, B
magnetic field, x spatial coordinate, v velocity
coordinate, and t time. The Vlasov operator d/dt describes the
streaming of particles influenced by electromagnetic forces, the
Maxwell's equations provide the electromagnetic fields, and the Landau
collision integral <cit.>, C, dissipates entropy and embodies the
transition from many-body dynamics to single particle statistics. As
such, the Vlasov-Maxwell-Landau system-of-equations is the gold
standard for high-fidelity fusion plasma simulations. 









The Vlasov-Maxwell-Landau system also conserves energy and momentum,
and guaranteeing these properties in numerical simulations is critical
to avoid plasma self heating and false momentum transfer during long-time simulations. Hirvijoki and Adams
recently developed a finite element discretization of the Landau
integral, which is able to preserve the conservation properties of the
Landau collision integral with sufficient order accurate finite
element space <cit.>.  We now continue this work with
the development of a multi-species Landau solver with adaptive mesh
refinement (AMR), which is designed for emerging architectures and implemented on the Second Generation Intel Xeon Phi,
Knights Landing (KNL), processor.

Due to the nonlinearity of the Landau collision integral, it has an intensive work complexity
of N^2 with N global integration or quadrature points.
Given this high-order work complexity, reducing the total number of
quadrature points decreases computational cost substantially.  We use
high order accurate finite elements and AMR to maximize the
information content of each quadrature point and thus minimize the
solver cost.  We adapt nonconforming tensor product meshes using the
p4est library
<cit.>,
as a third party library in the PETSc library
<cit.>.

We develop algorithms and techniques for optimizing the Landau solver
on emerging architectures, with emphasis on KNL, and verify the solver
on a model problem.  We vectorize the kernel using Intel AVX-512
intrinsics and achieve a flop rate as high as 22% of the
theoretical peak floating point rate of KNL.



§ CONSERVATIVE FINITE ELEMENT DISCRETIZATION OF THE LANDAU INTEGRAL



We consider the multi-species version of the conservative finite
element discretization of the Landau collision integral presented by
Hirvijoki and Adams <cit.>.  Under small-angle dominated
Coulomb collision, the distribution function f_α(v,t) of
species α evolves according to

    ∂ f_α/∂ t= ∑_βν_αβm_o/m_α∇_v·∫_Ω̅ dv̅ 𝐔(v,v̅)·(m_o/m_αf̅_β∇_v f_α - m_o/m_βf_α∇̅_v̅f̅_β).

Here
ν_αβ=e_α^2e_β^2lnΛ_αβ/(8π
m_o^2ε_0^2), lnΛ is the Coulomb logarithm, m_o
is an arbitrary reference mass, ε_0 is the vacuum
permittivity, m is mass, e is electric charge, and v is the
velocity. Overbar terms are evaluated on the v̅ grid that covers the domain Ω̅ of species β.  The
Landau tensor 𝐔(v,v̅) is a scaled projection
matrix defined as

    𝐔(v,v̅)=1/|v-v̅|^3(|v-v̅|^2𝐈-(v-v̅)(v-v̅))

and has an eigenvector v-v̅ corresponding to a zero
eigenvalue.

Given a test function ψ(v), the weak form of the Landau operator (<ref>) for species α is given by

    (ψ,∂ f_α/∂ t)_Ω=∑_β(ψ, f_α)_K,αβ+(ψ,f_α)_𝐃,αβ

where (·,·)_Ω is the standard L^2 inner product in Ω and the weighted inner products present the advective and diffusive parts of the Landau collision integral

    (ψ, ϕ)_K,αβ   =∫_Ωdv∇_vψ·ν̂_αβm_o/m_αm_o/m_βK(f_β,v) ϕ,
    (ψ,ϕ)_𝐃,αβ   =-∫_Ωdv∇_vψ·ν̂_αβm_o/m_αm_o/m_α𝐃(f_β,v)·∇_vϕ

The collision frequency is normalized with
ν̂_αβ=ν_αβ/ν_o so that time t is dimensionless, and f_β is
the distribution function of species β.  The vector
K and the tensor 𝐃 are defined as





    K(f,v)   =∫_Ω̅dv̅ 𝐔(v,v̅)·∇̅_v̅f(v̅), 
    𝐃(f,v)   =∫_Ω̅dv̅ 𝐔(v,v̅)f(v̅).











Assuming a finite-dimensional vector space V_h that is spanned by the set of functions {ψ_i}_i, the finite-dimensional approximation of the weak form (<ref>) can be written in a matrix form

    𝐌ḟ_α=𝐂_α[f]f_α

where f_α is the vector containing the projection coefficients of f_α onto V_h and the vector f is the collection of all species f_α, The mass and collision matrices are defined

    𝐌_ij= (ψ_i,ψ_j)_Ω,  𝐂_α,ij[f]=∑_β=1^S(ψ_i, ψ_j)_K,αβ+(ψ_i,ψ_j)_𝐃,αβ

The integrals in
(<ref>,<ref>), with the Landau tensors in the kernel, have
N work for each species β and each equation in
(<ref>).  With N equation this leads to an
N^2 algorithm for computing a Jacobian or residual when solving the equations (<ref>) for each species.




























§ ALGORITHM DESIGN FOR EMERGING ARCHITECTURES



This section discuses the algorithms and techniques used to
effectively utilize emerging architectures for a Landau integral
solver.  While the Landau operator has N^2 work complexity,
this work is amenable to vector processing.  We focus on KNL, but the
algorithm is design to minimized data movement and simplify access
patterns, which is beneficial for any emerging architecture. 

The discrete Landau Jacobian matrix construction, or residual
calculation, can be written as six nested loops.  Algorithm
<ref> shows high level pseudo-code for construction the
Landau Jacobian matrix, with |G| cells in the set G, Nq
quadrature points in each element, distribution functions f, S
species, and weights w_q_j = |J(q_j)| · q_j.weight
· q_j.r, where q_j.r is the axisymmetric term of the element
Jacobian, q_j.weight is the quadrature weight of q_j, and
J(q_j) is the element Jacobian at point q_j.


The Landau tensor 𝐔 in
(<ref>,<ref>) is computed, or read from memory, in the
inner loop.  
A vector K = 𝐔·∇ f_q_j w_q_j
and a tensor 𝐃 = 𝐔f_q_j w_q_j are
accumulated in the inner loop.  With S species, the accumulation of
K and 𝐃 requires 6S words.  These
accumulated values are transformed in a standard finite element
process from the reference to the real element geometry and assembled
with finite element shape functions into the element matrix.
The six loops of Algorithm <ref> can be processed in any
order, and blocked, giving different data access patterns, which is
critical in optimizing performance.

The first two issues that we address in the design of the Landau
solver are 1) whether to precompute the Landau tensors or compute them
as needed and 2) whether to use a single mesh with multiple degrees of
freedom per vertex or use a separate mesh for each species.





  §.§.§ To Precompute the Landau Tensor or Not to Precompute



The Landau tensor is only a function of mesh geometry and can be
computed and stored for each mesh configuration.  The cost of
computing the Landau tensor is amortized by the number of nonlinear
solver iterations and the number of time steps that the mesh is used
for.  The computation of the tensor can be expensive, especially in the axisymmetric case which involves two different tensors and evaluation of elliptic integrals (see Appendix <cit.>), requiring approximately 165 floating
point operations (flops) as measured by both the Intel Software
Development Emulator (SDE) and code analysis, including four logarithm
and square roots. Storing two such tensors requires eight words of storage,
64 bytes with double precision words.  There are N^2 unique
mesh (i,j) pairs for which the tensors are computed or stored, which is too much to fit in a
cache of any foreseeable machine with any reasonable degree of
accuracy (e.g., 64 megabytes with N=10^3).  The decision to
precompute or compute as needed depends on several factors.

A simple analysis on KNL suggest that both approaches are viable but
that trends in hardware works in favor of the compute as needed
approach.  Assuming the equivalent of 200 ordinary flops per axisymmetric Landau
tensor pair calculation and 64 bytes of data, the flop to byte ratio
is about three.  KNL has a theoretical peak floating point capacity of
about 2.6 × 10^12 flops/second and around 400 × 10^9
bytes/second on package memory bandwidth capacity, as measured by
STREAMS, or a flop to byte ratio of about six.  This simple analysis
suggests that precomputing would be two times slower, however, we
achieve about 20 % of theoretical peak flop rate and, thus, a
precomputed implementation would need only to achieve about 40% of
STREAMS bandwidth to match the run time of each kernel evaluation,
which one would expect is achievable.  This analysis suggests that
either method could be effective on KNL, but the spread between flop
capacity, in the form of more vector lanes and more hardware resources
per lane, and memory bandwidth capacity is anticipated to increase in
the future, which will benefit the compute on demand approach.

The kernel in Algorithms <ref> requires 3S+1 words from
memory per kernel evaluation for the weight w_q_j, the value
f_β( q_j ) and gradient of f_β( q_j
) for each species.  The compute on demand approach also
requires the coordinate.  This is N data, which has the
potential to fit in cache, for example, with N=10^3 and two species
this data would be about 64 kilobytes, plus lower order data, per
thread.  Eight threads per tile should fit in the 1Mb shared L2 cache
on KNL.



  §.§.§ Single and Multiple Meshes



We use a single mesh, independently adapted for each species, with S
degrees of freedom per vertex, however one can use multiple meshes or
a mesh for each species.  Observe that the integrals in
(<ref>,<ref>) are decoupled from the outer integral in
(<ref>,<ref>).  In theory, one can use a separate grid, or
different quadrature or even a different discretization, for each
species.  One could even use a different mesh for the inner and outer
integral in (<ref>,<ref>).  An advantage of using a single mesh
is that the two loops over species in (<ref>) can be
processed after the Landau tensor 𝐔
is computed, and hence this tensor can be reused S^2 times.
However, if all of the species have “orthogonal" optimal meshes, that
is each quadrature point only has significant information for one
species, which is a good assumption for ions and electrons because of
their disparate masses, then a single mesh requires about as many
vertices as the sum of each of the putative multiple meshes.  With the
model of orthogonal optimal meshes and kernel dominated computation or
communication, and with N_α quadrature points for each species
α, the complexity of a Landau solve is 𝒪(∑_α=1^S N_α)^2 for both the single and
multiple mesh approach.

With multiple ion species the orthogonal mesh assumption would be less
accurate, because (small) ions have similar optimal meshes, which
benefits the single mesh approach.  The Jacobian matrix for the single
mesh approach has about S times more non-zeros, which is not
important if the total cost is dominated by the Landau kernel.  It is
likely that with further optimization of the Landau kernel, the next
generations of hardware, and algorithmic improvements for the inner
integral, that the inner integral costs will be reduced relative to
the rest of the solver costs, which would benefit the multiple mesh
approach.  The single mesh method has larger accumulation register
demands and larger element matrices, which pressures the memory system
more and is advantageous for the multiple mesh approach.  The result
of the increased register pressure can be seen in the decreases flop
rates in Table <ref> with the increase in the number of
species.

Another potential advantage of the single grid method is that the
extra degrees of freedom, in for instance the ions in the range of one
electron thermal radius, might be beneficial.  The large scale
separation between ions and electrons means that small relative errors
in the ion distribution can be large relative to the electron
distribution.  Ions and electrons have about equal and opposite
charge, cancellation could lead to large relative errors in the total
charge density.
The accuracy of collisions with fast electrons in the tail of the ion
distribution could benefit from extra resolution.  A more thorough
understanding of the accuracy issues would be required to fully
address the question of using a single or multiple meshes.



 §.§ Our Algorithm



For demonstration purposes, we focus on implementing the axially
symmetric version using cylindrical velocity coordinates x=(r,θ,z). Under axial symmetry the distribution function is independent of the angular velocity coordinate (∂_θf=0) and the evaluation of the vector K and the tensor 𝐃 requires two different Landau tensors 𝐔_𝐊 and 𝐔_𝐃 respectively (see Appendix <cit.>). We choose to compute the
required Landau tensors as needed and use a single mesh with a
degree-of-freedom for each species on each vertex.  We fuse the two
inner loops over cells and quadrature points, inline the function call
of, and within, the Landau tensor function.  Algorithm <ref>
shows the initialization of the vectors r, z, w, f, and the
two gradient vectors df[1] and df[2], with |G| cells in the set
G, S species, and weights w_q_i at each quadrature point i.
Each quadrature point q_i is located at a 2D coordinate (q_i.r,
q_i.z).


Algorithm <ref> shows the algorithm for the construction of
the Landau collision integral Jacobian.

This algorithm is designed to minimized data movement by computing the
Landau tensors as needed and exploits a single mesh by lifting the
kernel outside of the two inner loops over species.





§ NUMERICAL METHODS AND IMPLEMENTATION



We implement the Landau solver with the PETSc numerical library
<cit.>.  PETSc provides finite element
(FE) and finite volume discretization support, mesh management,
interfaces to several third party mesh generators, fast multigrid
solvers, interfaces to several third party direct solvers, AMR
capabilities among other numerical methods.  We adapt nonconforming
tensor product meshes using the third party p4est library
<cit.>,
and unstructured conforming simplex meshes with PETSc's native AMR
capabilities <cit.>.  Our
experiments use bi-quadratic (Q2) elements with p4est
adaptivity, with the PETSc's Plex mesh management framework.

The computational domain is Ω={(r,z) | 0 ≤ r ≤ L, -L≤
z ≤ L}, with L=2.  We use Neumann boundary conditions and
shifted Maxwellian initial distribution functions, for each species,
of the form

    f_α(x,t=0)=θ1/2(πσ_α^2)^-3/2exp(-r^2+(z-s_α)^2/σ_α^2),

where σ_α=2T_αm_o/m_α, s_i = 0, s_e =
-1, T_α is temperature, and θ is a scaling factor
used to maintain quasi-neutrality.

We solve the boundary value problem 
    ∂
      f_α/∂ t( v, t ) -
    𝐂_α[f] f = 0
 
in axisymmetric coordinates, with standard FE methods and time
integrators.
A Newton nonlinear solver with the SuperLU direct liner solver is used
at each time stage or step <cit.>.  These
experiments use a Crank-Nicolson time integrator.

A global kinetic model would include a 3D spatial component and the 3V version of this
solver would be used at each cell in either a particle method
<cit.>, or a grid based kinetic method <cit.>.
Our numerical experiments use up to 272 Message Passing Interface
(MPI) processes on one KNL socket, redundantly solving the problem, to
include some of the memory contention of a full 5D code.  The timing
experiments run one time step with one Newton iteration, which results
in the Landau operator being called twice (one more than the number of
Newton iterations), and with one linear solve (one per Newton
iteration).  The time for this step is reported, which does not
include the AMR mesh construction.  We observe a significant
variability in times with 272 processes and have run each test several
times in several sessions, in both batch and interactive sessions, and
the report the fastest observed time.  We report the maximum time from any processor and see about a
10% ratio between the maximum and the minimum time of any process
with large process counts.



 §.§ Overview of application and numerical method



To illustrate the capabilities and behavior of the solver, we run the
code to near equilibrium, initializing electrons with a shifted
Maxwellian distribution hitting a stationary single proton ion
population with a Maxwellian distribution.  We use a realistic mass
ratio of m_i/m_e=1836.5 and temperatures of T_e = 0.02 and
T_i = 0.002.  Figure <ref> (left) shows the initial
electron distribution with the ion grid at the origin, a partially
thermalized electron distribution (center, left), and Maxwellian ion
distribution near equilibrium (center, right).  The ion distribution
has been shifted from the origin by collision with the electrons.


The ions are resolved with AMR at the origin and have a near
Maxwellian distribution.  Note, Figure <ref> uses linear
interpolation from the four corners of each quadrilateral, whereas the
numerics use bi-quadratic interpolation with nine vertices per
quadrilateral, which results in inaccuracy and asymmetries in the
plots not present in the numerics.



 §.§ Optimization and Performance



Most of the work in the Landau solver is in the inner integral of
(<ref>,<ref>) (lines 8-16 in Algorithm <ref>).  This
kernel is vectorized with Intel AVX-512 intrinsics.  The Landau
tensors calculation includes two logarithms, a square root, a power,
seven divides, about 85 multiplies and 165 total flops.  The power is
converted to a inverse square root, an intermediate divide is reused,
resulting in five divides, two logarithms, a square root, and an
inverse square root.  The KNL sockets used for this study is equipped
with 34 tiles, each with a 1 megabyte shared L2 cache and two cores,
each core has two 8 lane vector units that can issue one fused
multiply add (FMA) per cycle per lane.  Each core has four hardware
threads, for a total of 272 threads per socket.  The KNL clock rate is
1.4 GHz, but is clocked down to 1.2 GHz in AVX-512 code segments.
This results in a theoretical flop peak rate of 2.6 × 10^12
flops/second.  The peak flop rate that can be achieved with this
solver is reduced because the kernel is not entirely composed of FMAs
and the four logarithms and square roots require considerably more
than one cycle each.

The performance data in this section uses a simplified version the
test problem, a grid adapted for electrons only with 176 cells and
1,584 quadrature points, a mass ratio of m_i/m_e=1 and T_e
= T_i = 0.2, and no Maxwellian shifts (s_i = s_e = 0), as shown in
Figure <ref> (right).



 §.§ Performance overview



The major code segments have been instrumented with PETSc timers.
Table <ref> show the maximum time from any process for major components of the Landau operator, the total Landau operator, and the
linear solver.


This data is from the double precision solve with 272 processes and
two species in Table <ref>.  This data shows that the
Landau kernel, though vectorized, is still responsible for most of the run time.



 §.§ Performance and complexity analysis



There are two types of work in the kernel: 1) computing the two Landau
tensors and 2) the accumulation of K vector and
𝐃 tensor.  The accumulation requires 20S^2 flops
(lines 12-13 in Algorithm <ref>).  Instrumenting this inner
loop would be invasive, but we can infer the percentage of time and
work in these two parts with a complexity model and global
measurements.  Assume both the time and work cost of the entire solver
are of the form C = aS^0 + bS^1 + cS^2.  The solve time and flop
counts with S=1,2,3, shown in Table <ref>, generate
right hand sides for system of three equations and three unknows a, b, and c, which are the time spent, or work, in each
of the three types of components.

The Landau tensor cost is formally independent of the number of
species, and the work in the accumulation has S^2 work
complexity.  Most of the rest of the costs, given a mesh, order of
elements, etc., has S^1 complexity.  Table <ref>
shows the percentage of time and work in each component, infered from
the data in Table <ref>.

This analysis shows that about 98% of the work, and about 88% of
the time, is in the kernel with one process.  The measurements of the
kernel time in Table <ref> is 66% with 272 processes.
This discrepancy is probably due to performance noise and memory
contention in the 272 process timings.  The non-kernel time percentage
(12) increases by a factor of about eight from the flop percentage
(1.5), which reflects the eight vector lanes of the KNL vector unit.



 §.§ Memory performance


Our experiments are trivially parallel, but memory contention results
in performance degradation as more processes are used on a socket.
KNL's architecture allows for twice as many vector lanes with single
(32 bit) versus double (64 bit) precision and can thus run, in theory
twice as fast in single precision.  This section investigates memory
issues with a weak speedup study in single and double precision.
Table <ref> shows timing data with increasing number of
processes on a single KNL socket, with single and double precision.

This data shows that we are achieving about 80% of the perfect
factor of two speedup with single precision.

A KNL socket has 136 vector units.  One might expect that using more
processes than 136 would not be useful, however, the kernel has serial
dependancies that result in bubbles in the pipeline, especially in the
ninth and tenth order polynomial evaluations in the elliptic
integrals.  These holes can be filled by interleaving a second process
in the second hardware thread.  We do see about a 15% increase in
total performance from the added parallelism of using 272 processes.
This data shows that the flop rate per process decreases by a factor
of about 3 in going from one process per tile to two processes per
vector unit, or eight processes per tile.  There is no difference
between the one process and 34 process runs, suggesting that the 34
processes are indeed placed with one per tile.  There is little
degradation going from one to two processes per tile, suggesting that
the problem still fits in the L2 cache.  The degradation from one to
two processes per core is from L1 cache misses because four processes
should be able to fit into the 1 megabyte L2 cache.  Analysis of
memory complexity and the time data going from one to two processes
per vector unit suggest that the full 272 process runs fit in the L2
cache, or nearly so.



 §.§ Verification


Consider a convergence study with Cartesian grids of the third moment,
thermal flux, to verify the expected order of convergence.  An
analytical flux for this problem does not exist.  Richardson
extrapolation is used to construct an approximate exact flux.  The
mass ratio is 4, T_e = 0.2 and T_i = 0.02.  The flux history,
with a series of refined grids starting with a 8 × 16 grid, is
shown in Figure <ref> (top, left).  Figure <ref>
also shows the differences between fluxes on successive grids, and the
error convergence.


We can see from this data that we achieve fourth order convergence.




§ CLOSURE


We have implemented a high order accurate finite element
implementation of the Landau collision integral with adaptive mesh
refinement in the PETSc library using AVX-512 intrinsics for the
Second Generation Intel Xeon Phi, Knights Landing, processor.  We have
developed a memory centric algorithm for emerging
architectures that is amenable to vector processing.  We have achieved up to 22% of
the theoretical peak flop rate of KNL and analyzed the performance
characteristics of the algorithm with respect to process memory
contention, single and double precision, and the results of
vectorization.
We have verified fourth order accuracy with a
bi-quadratic, Q2, finite element discretization.
Future work includes, building models for runaway electrons
in tokamak plasmas with this kernel <cit.>, 
and building up complete kinetic models (6D AMR) that also preserve the geometric 
structure of the governing equations of fusion plasmas <cit.>.




§ ACKNOWLEDGMENTS


We are grateful for discussions and insights from Intel engineer Vamsi
Sripathi.  This work has benefited from many discussions with Sam
Williams.  This work was partially funded from
the Intel Parallel Computing Center program.  This research used
resources of the National Energy Research Scientific Computing Center,
a DOE Office of Science User Facility supported by the Office of
Science of the U.S. Department of Energy under Contract
No. DE-AC02-05CH11231.

siamplain


