Timings and parallelization

This section summarizes the options DFTK offers to monitor and influence performance of the code. For details on running DFTK on HPC systems, see also Using DFTK on compute clusters and Using DFTK on GPUs.

Timing measurements

By default DFTK uses TimerOutputs.jl to record timings, memory allocations and the number of calls for selected routines inside the code. These numbers are accessible in the object DFTK.timer. Since the timings are automatically accumulated inside this datastructure, any timing measurement should first reset this timer before running the calculation of interest.

For example to measure the timing of an SCF:

DFTK.reset_timer!(DFTK.timer)
scfres = self_consistent_field(basis, tol=1e-5)

DFTK.timer

────────────────────────────────────────────────────────────────────────────────
                                       Time                    Allocations      
                              ───────────────────────   ────────────────────────
      Tot / % measured:            331ms /  44.5%           98.6MiB /  66.8%    

Section               ncalls     time    %tot     avg     alloc    %tot      avg
────────────────────────────────────────────────────────────────────────────────
self_consistent_field      1    147ms  100.0%   147ms   65.9MiB  100.0%  65.9MiB
  LOBPCG                  33   75.7ms   51.5%  2.29ms   14.4MiB   21.8%   446KiB
    DftHamiltonian...     88   65.3ms   44.4%   742μs   5.15MiB    7.8%  59.9KiB
      local              580   63.1ms   42.9%   109μs    118KiB    0.2%     208B
      nonlocal            88    949μs    0.6%  10.8μs    373KiB    0.6%  4.24KiB
    ortho! X vs Y         77   3.28ms    2.2%  42.6μs   1.35MiB    2.0%  17.9KiB
      ortho!             167   1.91ms    1.3%  11.4μs   1.01MiB    1.5%  6.18KiB
    rayleigh_ritz         55   3.07ms    2.1%  55.7μs   1.24MiB    1.9%  23.1KiB
    ortho!                33    499μs    0.3%  15.1μs    166KiB    0.2%  5.03KiB
    preconditioning       88    321μs    0.2%  3.65μs   24.0KiB    0.0%     279B
    Update residuals      88    235μs    0.2%  2.67μs   42.4KiB    0.1%     493B
  compute_density         11   35.8ms   24.4%  3.26ms   7.04MiB   10.7%   655KiB
    symmetrize_ρ          11   29.0ms   19.7%  2.63ms   5.44MiB    8.3%   507KiB
  energy_hamiltonian      12   17.1ms   11.6%  1.42ms   18.2MiB   27.6%  1.52MiB
    ene_ops               12   15.3ms   10.4%  1.28ms   13.6MiB   20.6%  1.13MiB
      ene_ops: xc         12   12.6ms    8.6%  1.05ms   6.66MiB   10.1%   568KiB
      ene_ops: har...     12   1.87ms    1.3%   155μs   5.41MiB    8.2%   462KiB
      ene_ops: non...     12    287μs    0.2%  23.9μs    181KiB    0.3%  15.1KiB
      ene_ops: kin...     12    185μs    0.1%  15.4μs    116KiB    0.2%  9.64KiB
      ene_ops: local      12    177μs    0.1%  14.7μs   1.10MiB    1.7%  93.7KiB
  energy                  11   11.3ms    7.7%  1.02ms   10.4MiB   15.8%   971KiB
    energy: xc            11   8.81ms    6.0%   801μs   4.06MiB    6.2%   378KiB
    ene_ops: hartree      11   1.60ms    1.1%   145μs   4.96MiB    7.5%   462KiB
    ene_ops: nonlocal     11    306μs    0.2%  27.8μs    181KiB    0.3%  16.5KiB
    ene_ops: kinetic      11    185μs    0.1%  16.8μs    107KiB    0.2%  9.77KiB
    ene_ops: local        11    175μs    0.1%  15.9μs   1.01MiB    1.5%  93.7KiB
  Anderson acceler...     10   3.60ms    2.5%   360μs   7.11MiB   10.8%   728KiB
  ortho_qr                 3    148μs    0.1%  49.5μs    101KiB    0.1%  33.7KiB
  χ0Mixing                11   75.0μs    0.1%  6.82μs   37.1KiB    0.1%  3.38KiB
enforce_real!              1   56.0μs    0.0%  56.0μs      384B    0.0%     384B
────────────────────────────────────────────────────────────────────────────────

The output produced when printing or displaying the DFTK.timer now shows a nice table summarising total time and allocations as well as a breakdown over individual routines.

Rough timing estimates

A very (very) rough estimate of the time per SCF step (in seconds) can be obtained with the following function. The function assumes that FFTs are the limiting operation and that no parallelisation is employed.

function estimate_time_per_scf_step(basis::PlaneWaveBasis)
    # Super rough figure from various tests on cluster, laptops, ... on a 128^3 FFT grid.
    time_per_FFT_per_grid_point = 30 #= ms =# / 1000 / 128^3

    (time_per_FFT_per_grid_point
     * prod(basis.fft_size)
     * length(basis.kpoints)
     * div(basis.model.n_electrons, DFTK.filled_occupation(basis.model), RoundUp)
     * 8  # mean number of FFT steps per state per k-point per iteration
     )
end

"Time per SCF (s):  $(estimate_time_per_scf_step(basis))"

"Time per SCF (s):  0.008009033203124998"

Options for parallelization

At the moment DFTK offers two ways to parallelize a calculation, firstly shared-memory parallelism using threading and secondly multiprocessing using MPI (via the MPI.jl Julia interface). MPI-based parallelism is currently only over $k$-points, such that it cannot be used for calculations with only a single $k$-point. There is also support for Using DFTK on GPUs.

The scaling of both forms of parallelism for a number of test cases is demonstrated in the following figure. These values were obtained using DFTK version 0.1.17 and Julia 1.6 and the precise scalings will likely be different depending on architecture, DFTK or Julia version. The rough trends should, however, be similar.

The MPI-based parallelization strategy clearly shows a superior scaling and should be preferred if available.

MPI-based parallelism

Currently DFTK uses MPI to distribute on $k$-points only. This implies that calculations with only a single $k$-point cannot use make use of this. For details on setting up and configuring MPI with Julia see the MPI.jl documentation.

1. First disable all threading inside DFTK, by adding the following to your script running the DFTK calculation:

   using DFTK
   disable_threading()

2. Run Julia in parallel using the mpiexecjl wrapper script from MPI.jl:

   mpiexecjl -np 16 julia myscript.jl

   In this -np 16 tells MPI to use 16 processes and -t 1 tells Julia to use one thread only. Notice that we use mpiexecjl to automatically select the mpiexec compatible with the MPI version used by MPI.jl.

As usual with MPI printing will be garbled. You can use

DFTK.mpi_master() || (redirect_stdout(); redirect_stderr())

at the top of your script to disable printing on all processes but one.

Thread-based parallelism

Threading in DFTK currently happens on multiple layers distributing the workload over different $k$-points, bands or within an FFT or BLAS call between threads. At its current stage our scaling for thread-based parallelism is worse compared MPI-based and therefore the parallelism described here should only be used if no other option exists. To use thread-based parallelism proceed as follows:

1. Ensure that threading is properly setup inside DFTK by adding to the script running the DFTK calculation:

   using DFTK
   setup_threading()

   This disables FFT threading and sets the number of BLAS threads to the number of Julia threads.

2. Run Julia passing the desired number of threads using the flag -t:

   julia -t 8 myscript.jl

For some cases (e.g. a single $k$-point, fewish bands and a large FFT grid) it can be advantageous to add threading inside the FFTs as well. One example is the Caffeine calculation in the above scaling plot. In order to do so just call setup_threading(n_fft=2), which will select two FFT threads. More than two FFT threads is rarely useful.

Advanced threading tweaks

The default threading setup done by setup_threading is to select one FFT thread, and use Threads.nthreads() to determine the number of DFTK and BLAS threads. This section provides some info in case you want to change these defaults.

BLAS threads

All BLAS calls in Julia go through a parallelized OpenBlas or MKL (with MKL.jl. Generally threading in BLAS calls is far from optimal and the default settings can be pretty bad. For example for CPUs with hyper threading enabled, the default number of threads seems to equal the number of virtual cores. Still, BLAS calls typically take second place in terms of the share of runtime they make up (between 10% and 20%). Of note many of these do not take place on matrices of the size of the full FFT grid, but rather only in a subspace (e.g. orthogonalization, Rayleigh-Ritz, ...) such that parallelization is either anyway disabled by the BLAS library or not very effective. To set the number of BLAS threads use

using LinearAlgebra
BLAS.set_num_threads(N)

where N is the number of threads you desire. To check the number of BLAS threads currently used, you can use BLAS.get_num_threads().

DFTK threads

On top of BLAS threading DFTK uses Julia threads in a couple of places to parallelize over $k$-points (density computation) or bands (Hamiltonian application). The number of threads used for these aspects is controlled by default by Threads.nthreads(), which can be set using the flag -t passed to Julia or the environment variable JULIA_NUM_THREADS. Optionally, you can use setup_threading(; n_DFTK) to set an upper bound to the number of threads used by DFTK, regardless of Threads.nthreads() (for instance, if you want to thread several DFTK runs and don't want DFTK's threading to interfere with that).

FFT threads

Since FFT threading is only used in DFTK inside the regions already parallelized by Julia threads, setting FFT threads to something larger than 1 is rarely useful if a sensible number of Julia threads has been chosen. Still, to explicitly set the FFT threads use

using FFTW
FFTW.set_num_threads(N)

where N is the number of threads you desire. By default no FFT threads are used, which is almost always the best choice.