How to distribute a calculation over MPI ranks?

How to distribute a calculation over MPI ranks?#

When every MPI rank holds the same structure, each rank would normally repeat the complete dispersion calculation. A work partition assigns a disjoint share of the pairwise, three-body and reciprocal space loops to each rank instead.

Parts are zero based and every unit of work belongs to exactly one part, so summing the energy, gradient and virial over all parts reproduces the complete result. D3 itself performs no communication, the reduction is left to the caller.

Note

The \(C_6\) coefficients are only evaluated for the pairs a part consumes, unless the three-body term is active, which reads them for all pairs. The coordination number is consumed in full by every part and stays unpartitioned unless the parts can exchange it, see Partition the coordination number.

To test the examples you can install the dependencies with

Fortran

mamba create d3 simple-dftd3 fortran-compiler openmpi pkg-config
mamba activate d3

C

mamba create d3 simple-dftd3 c-compiler openmpi pkg-config
mamba activate d3

Partition the interaction loops#

The Fortran API takes the partition as an optional argument, omitting it selects the complete work. The C API stores it on the dispersion model, next to the cutoffs, and so does the Python API with DispersionModel.set_work_partition.

Fortran

mpi.f90#

program mpi_partition
   use, intrinsic :: iso_fortran_env, only : r8 => real64
   use dftd3, only: d3_model, d3_param, rational_damping_param, get_rational_damping, &
      & new_rational_damping, new_d3_model, get_dispersion, realspace_cutoff, &
      & work_partition, new_work_partition
   use mctc_env, only: error_type
   use mctc_io, only: structure_type, new
   use mpi_f08, only : MPI_COMM_WORLD, MPI_DOUBLE_PRECISION, MPI_IN_PLACE, MPI_SUM, &
      & MPI_Abort, MPI_Finalize, MPI_Init, MPI_Comm_rank, MPI_Comm_size, &
      & MPI_Allreduce
   implicit none

   character(len=:), allocatable :: method
   type(structure_type) :: mol
   type(error_type), allocatable :: error
   integer, allocatable :: num(:)
   real(r8), allocatable :: xyz(:, :)
   real(r8), allocatable :: gradient(:, :)
   real(r8) :: energy

   type(d3_model) :: disp
   type(d3_param) :: inp
   type(rational_damping_param) :: param
   type(work_partition) :: partition
   integer :: rank, nranks

   call MPI_Init()
   call MPI_Comm_rank(MPI_COMM_WORLD, rank)
   call MPI_Comm_size(MPI_COMM_WORLD, nranks)

   method = "PBE0"
   num = [6, 1, 1, 1, 1]
   xyz = reshape([ &  ! coordinates in Bohr
     &  0.0000000_r8, -0.0000000_r8,  0.0000000_r8, &
     & -1.1922080_r8,  1.1922080_r8,  1.1922080_r8, &
     &  1.1922080_r8, -1.1922080_r8,  1.1922080_r8, &
     & -1.1922080_r8, -1.1922080_r8, -1.1922080_r8, &
     &  1.1922080_r8,  1.1922080_r8, -1.1922080_r8],&
     & [3, size(num)])
   call new(mol, num, xyz, charge=0.0_r8, uhf=0)

   call get_rational_damping(inp, method, error, s9=1.0_r8)
   if (allocated(error)) call fatal(error%message)
   call new_rational_damping(param, inp)
   call new_d3_model(disp, mol)

   ! every rank holds the complete structure and evaluates its own share of the
   ! pairwise and three-body interactions, parts are zero based
   call new_work_partition(error, partition, rank, nranks)
   if (allocated(error)) call fatal(error%message)

   allocate(gradient(3, mol%nat))
   call get_dispersion(error, mol, disp, param, realspace_cutoff(), energy, gradient, &
      & partition=partition)
   if (allocated(error)) call fatal(error%message)

   ! s-dftd3 performs no communication, the caller reduces the partial results
   call MPI_Allreduce(MPI_IN_PLACE, energy, 1, MPI_DOUBLE_PRECISION, MPI_SUM, MPI_COMM_WORLD)
   call MPI_Allreduce(MPI_IN_PLACE, gradient, size(gradient), MPI_DOUBLE_PRECISION, &
      & MPI_SUM, MPI_COMM_WORLD)

   if (rank == 0) then
      print "(3a, f13.10, a)", "Dispersion energy for ", method, "-D3(BJ) is ", energy, " Hartree"
   end if

   call MPI_Finalize()

contains

   subroutine fatal(message)
      character(len=*), intent(in) :: message
      print "(2a)", "Error: ", message
      call MPI_Abort(MPI_COMM_WORLD, 1)
   end subroutine fatal

end program mpi_partition

C

mpi.c#

#include <stdio.h>
#include <stdbool.h>

#include <mpi.h>

#include "dftd3.h"

int main(int argc, char **argv)
{
  dftd3_error error = dftd3_new_error();
  dftd3_structure mol = NULL;
  dftd3_model d3 = NULL;
  dftd3_param param = NULL;
  int rank, nranks;

  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &nranks);

  int nat = 5;
  int num[5] = {6, 1, 1, 1, 1};
  double xyz[15] = {  // coordinates in Bohr
     0.00000000, -0.00000000,  0.00000000,
    -1.19220800,  1.19220800,  1.19220800,
     1.19220800, -1.19220800,  1.19220800,
    -1.19220800, -1.19220800, -1.19220800,
     1.19220800,  1.19220800, -1.19220800};
  mol = dftd3_new_structure(error, nat, num, xyz, NULL, NULL);
  if (dftd3_check_error(error)) goto handle_error;

  char method[5] = "PBE0";
  param = dftd3_load_rational_damping(error, method, true);
  if (dftd3_check_error(error)) goto handle_error;

  d3 = dftd3_new_d3_model(error, mol);
  if (dftd3_check_error(error)) goto handle_error;

  // every rank holds the complete structure and evaluates its own share of the
  // pairwise and three-body interactions, parts are zero based
  dftd3_set_model_work_partition(error, d3, rank, nranks);
  if (dftd3_check_error(error)) goto handle_error;

  double energy;
  double gradient[15];
  dftd3_get_dispersion(error, mol, d3, param, &energy, gradient, NULL);
  if (dftd3_check_error(error)) goto handle_error;

  // s-dftd3 performs no communication, the caller reduces the partial results
  MPI_Allreduce(MPI_IN_PLACE, &energy, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  MPI_Allreduce(MPI_IN_PLACE, gradient, 3 * nat, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

  if (rank == 0) {
    printf("Dispersion energy for %s-D3(BJ) is %13.10lf Hartree\n", method, energy);
  }

  dftd3_delete(error);
  dftd3_delete(mol);
  dftd3_delete(d3);
  dftd3_delete(param);
  MPI_Finalize();
  return 0;

handle_error: {
  char msg[512];
  dftd3_get_error(error, msg, NULL);
  printf("Error: %s\n", msg);

  dftd3_delete(error);
  dftd3_delete(mol);
  dftd3_delete(d3);
  dftd3_delete(param);
  MPI_Abort(MPI_COMM_WORLD, 1);
  return 1;
}
}

Fortran

❯ mpifort mpi.f90 $(pkg-config s-dftd3 mctc-lib --cflags --libs) && mpirun -n 4 ./a.out
Dispersion energy for PBE0-D3(BJ) is -0.0009218696 Hartree

C

❯ mpicc mpi.c $(pkg-config s-dftd3 --cflags --libs) && mpirun -n 4 ./a.out
Dispersion energy for PBE0-D3(BJ) is -0.0009218696 Hartree

The result is independent of the number of ranks up to the summation order. Running with a single rank is identical to omitting the partition, which is also available as serial_work_partition if you prefer to always pass it explicitly.

The analytical Hessian and the geometric counterpoise correction accept the same partition, the pairwise decomposition does not.

Let D3 do the reduction#

MPI support is an opt-in build option:

meson

meson setup _build -Dmpi=true

CMake

cmake -B _build -DSDFTD3_WITH_MPI=ON

It enables the dftd3_mpi module, which derives the partition from a communicator and reduces the results, replacing the whole boilerplate above. The C API takes the communicator on the model instead of a part index, everything else stays the same:

Fortran

use dftd3, only : get_dispersion_mpi
use mpi, only : MPI_COMM_WORLD

call get_dispersion_mpi(error, mol, disp, param, realspace_cutoff(), &
   & MPI_COMM_WORLD%mpi_val, energy, gradient)

C

dftd3_set_model_mpi_comm(error, disp, MPI_Comm_c2f(MPI_COMM_WORLD));
dftd3_get_dispersion(error, mol, disp, param, &energy, gradient, NULL);

Communicators are passed as plain Fortran handles, users of mpi_f08 pass comm%mpi_val and C callers convert with MPI_Comm_c2f. The library itself is built against mpi_f08. The counter-poise correction has the same entry points, get_counterpoise_mpi and dftd3_set_gcp_mpi_comm. For everything else, new_mpi_work_partition creates the partition of the calling rank and leaves the reduction to you.

Note

The routines are collective, every rank of the communicator has to call them with the same input.

Without MPI support the entry points are still there, but every one of them reports the missing feature in the error handler instead. The build configuration can be queried at compile time and at runtime:

Fortran

use dftd3, only : dftd3_has_mpi, dftd3_has_feature

if (dftd3_has_feature("mpi")) then
   ! ...
end if

C

if (dftd3_has_feature("mpi")) {
   /* ... */
}

The command line driver reports the same in s-dftd3 --version.

Warning

The command line driver itself is a serial program, it never initializes MPI and only reports what the library it linked against supports. Launching it with mpirun -n 4 s-dftd3 ... starts four independent complete calculations writing to the same output files, it does not distribute one. Distributing a calculation is done from your own program through the APIs above.

Note

The Python API binds the work partition but not the communicator entry points, since taking a communicator would make mpi4py a dependency of the package. Combining set_work_partition with your own mpi4py reduction works and gives the same result, only the coordination number stays unpartitioned there. This can be revisited if there is demand.

Partition the coordination number#

The coordination number is needed in full by every part, so it can only be partitioned if the parts exchange it halfway through the calculation. Passing a work_reducer alongside the partition enables this, get_dispersion_mpi does so automatically. Bring your own communication layer by implementing the deferred reduce binding:

type, extends(work_reducer) :: my_reducer
contains
   procedure :: reduce => my_reduce
end type my_reducer

! sum val over all parts, in place
subroutine my_reduce(self, val, error)
   class(my_reducer), intent(in) :: self
   real(wp), intent(inout), contiguous :: val(:)
   type(error_type), allocatable, intent(out) :: error
end subroutine my_reduce

The reducer is called twice per gradient evaluation, both times on an array of the size of the system. Without it the coordination number stages stay unpartitioned and the results are unchanged.

How to distribute a calculation over MPI ranks?

Contents

How to distribute a calculation over MPI ranks?#

Partition the interaction loops#

Let D3 do the reduction#

Partition the coordination number#