Multithreading and MPI#

As typical operations on time-ordered data can be quite consuming, the LiteBIRD Simulation Framework provides a number of tools to exploit the presence of multiple CPU cores and even multiple computing nodes. This section details how to take advantage of these facilities and is split in two parts:

  • We will first present the ability of the framework to use multiple CPU threads; in this context, the data samples are kept in a chunk of memory that is shared between several processes. The framework uses Numba, which can take advantage either of the Intel Threading Building Blocks library or of OpenMP.

  • Then, we will discuss the possibility to run the code on multiple computing nodes, where the memory of each node is not shared with the others. The framework is able to use any MPI library, through the Python package mpi4py.

Multithreading#

Some parts of the LiteBIRD Simulation Framework are able to exploit multiple cores because several of its modules rely on the Numba library.

To use more than one CPU, you must explicitly tell the Framework how many to use, either by setting the environment variable OMP_NUM_THREADS to the number of CPUs to use, or by using two parameters in the constructor of the class Simulation. The Framework does not default to using every CPU core available on the machine: doing so would be unsafe under MPI, where several ranks typically share a node and would otherwise all try to grab every core at once. Leaving thread counts unconfigured makes the Framework run single-threaded, both for Numba and for ducc0 (see Threads and ducc0 below); this is a safe default, but it means that on a single-user laptop or workstation you do need to set OMP_NUM_THREADS (or the equivalent parameters below) yourself if you want to take advantage of all the available cores.

  • numba_threads: this is the number of CPUs that Numba will use for parallel calculations. The parameter defaults to None, in which case it is resolved once, following this order of precedence:

    1. the value of numba_threads itself, if you passed one explicitly (or if it was set in a TOML parameter file, see below);

    2. the environment variable NUMBA_NUM_THREADS, if set;

    3. the environment variable OMP_NUM_THREADS, if set;

    4. 1, if none of the above is set.

    This is the same resolution order used for the nthreads parameters accepted by the low-level ducc0-based functions of the framework (spherical harmonic transforms, beam convolution, map scanning, …) — see Threads and ducc0 below. Using NUMBA_NUM_THREADS instead of OMP_NUM_THREADS lets you give Numba a different thread count than ducc0, if you ever need to; leaving NUMBA_NUM_THREADS unset makes the two agree by default.

  • numba_threading_layer: this parameter is a string that specifies which threading library should be used by Numba. The value depends both on the version of Numba you are running and on the availability of these libraries, as they are not installed together with the LiteBIRD Simulation Framework. Numba 0.53 provides the following choices:

    • tbb: Intel Threading Building Blocks. You should pick this if you are running your code on Intel machines and the Tbb library is available.

    • omp: OpenMP. If you pick this one, be sure that the OpenPM library is available.

    • workqueue: this is an internal threading library provided by Numba. It’s probably the least efficient of the three; its main advantage is that it is always available.

These parameters can be passed through a TOML parameter file (see Parameter files) as well:

# This is file "my_conf.toml"
[simulation]
random_seed = 12345
numba_threads = 32
numba_threading_layer = "tbb"

Both tbb and omp require that the relevant library be available on your system, as the command pip install litebird_sim does not install them. If you are running your code on a HPC cluster, it is probably a matter of running a command like the following:

# This might change depending on how the environment on your cluster
# is configured; the following commands are just examples.
$ module load tbb     # Intel Threading Building Blocks
$ module load openmp  # OpenMP

Threads and ducc0#

Besides Numba, most of the computationally-heavy, low-level functions of the framework (spherical harmonic transforms in litebird_sim.maps_and_harmonics, beam convolution, map scanning, the map-makers, …) delegate their parallel work to the ducc0 library through an nthreads parameter. Every such function defaults to nthreads=None, which is resolved through litebird_sim.resolve_nthreads() using the same precedence as OMP_NUM_THREADS for Numba above: explicit value, then OMP_NUM_THREADS, then 1. The values used by the current process, resolved once when litebird_sim is imported, are exposed as litebird_sim.NUM_THREADS (for ducc0) and litebird_sim.NUMBA_NUM_THREADS (for Numba).

This 1-thread fallback is deliberately conservative: it keeps both serial and MPI runs safe by default (no risk of a rank grabbing every core on a shared node), at the cost of not using extra cores unless you ask for them. If you want multithreaded ducc0/Numba execution — whether on a laptop or per MPI rank on a cluster — set OMP_NUM_THREADS (and, if needed, NUMBA_NUM_THREADS) to the number of cores you want to use before launching your job, e.g. for 8 MPI ranks on a 64-core node, using 8 threads per rank:

$ export OMP_NUM_THREADS=8
$ mpirun -n 8 python3 my_script.py

You can always override the resolved value for a single call by passing an explicit nthreads argument to the function you are calling.

MPI#

The LiteBIRD Simulation Framework lists mpi4py as an optional dependency. This means that simulation codes should be able to cope with the lack of MPI.

The framework can be forced to use MPI or not using the variable LITEBIRD_SIM_MPI:

  • Set the variable to 1 or an empty value to force importing mpi4py;

  • Set the variable to 0 to avoid importing mpi4py;

  • If the variable is not set, the code will try to import mpi4py, but in case of error it will not complain and will silently shift to serial execution.

The framework provides a global variable, MPI_COMM_WORLD, which is the same as mpi4py.MPI.COMM_WORLD if MPI is being used. Otherwise, if MPI is not being used, it still contains the following members:

  • rank (set to 0);

  • size (set to 1).

Thus, the following code works regardless whether MPI is present or not:

import litebird_sim as lbs

if lbs.MPI_COMM_WORLD.rank == 0:
    print("Hello, world!")

However, you can use MPI_COMM_WORLD to call MPI functions only if MPI was actually enabled. You can check this using the Boolean variable MPI_ENABLED:

import litebird_sim as lbs

comm = lbs.MPI_COMM_WORLD
if lbs.MPI_ENABLED:
    comm.barrier()

To ensure that your code uses MPI in the proper way, you should always use MPI_COMM_WORLD instead of importing mpi4py directly.

Enabling/disabling MPI#

The user can control whether MPI must be used or not in a script, through the environment variable LITEBIRD_SIM_MPI (ENABLE_MPI is accepted as well):

  • If the variable is set to the empty string or to 1, true, on, yes, then mpi4py is imported, and an exception is raised if this cannot be done (e.g., because it was not installed using the flag --extra mpi when uv sync was called).

  • If the variable is set to 0, false, off or no, then mpi4py is not imported, even if it is installed.

  • If the variable is not set, then mpi4py will be imported, but any failure will be accepted and the framework will silently switch to serial mode.

Grasping how MPI is being used#

You will typically use MPI to spread TODs among many MPI processes, so that the simulation can span several detectors and a longer time scale. Unfortunately, this means that it’s often complicated to understand how data is being kept in memory.

If you use the Simulation object (and you should, you really should!), you can call Simulation.describe_mpi_distribution() after you have allocated the TODs via Simulation.create_observations(); it will return an instance of the class MpiDistributionDescr, which can be inspected and printed to the terminal. See Section Simulations for more information about this.

API reference#

litebird_sim.mpi.MPI_COMM_WORLD = <litebird_sim.mpi._SerialMpiCommunicator object>#

Global variable equal either to mpi4py.MPI.COMM_WORLD or a object that defines the member variables rank = 0 and size = 1.

litebird_sim.mpi.MPI_CONFIGURATION = {}#

If MPI_ENABLED is True, this is a dictionary containing information about the MPI configuration. Otherwise, it is an empty dictionary

litebird_sim.mpi.MPI_ENABLED = False#

True if MPI should be used by the application. The value of this variable is set according to the following rules:

  • If the environment variable LITEBIRD_SIM_MPI is set to 1, use MPI and fail if mpi4py cannot be imported;

  • If the environment variable LITEBIRD_SIM_MPI is set to 0, avoid using MPI even if mpi4py is present;

  • If the environment variable LITEBIRD_SIM_MPI is not set, try to use MPI and gracefully revert to a serial mode of execution if mpi4py cannot be imported.