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 and ATM interaction 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. DFT-D4 itself performs no communication; the reduction is left to the caller.

Note

Structure-dependent quantities such as coordination numbers, charges and \(C_6\) coefficients are evaluated for the full system on every part. The speedup is therefore bounded by the interaction loops, which dominate for larger systems and when the ATM contribution is enabled.

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.

Fortran

mpi.f90

program mpi_partition
   use, intrinsic :: iso_fortran_env, only : r8 => real64
   use dftd4, only : d4_model, damping_param, get_dispersion, &
      & get_rational_damping, new_d4_model, new_work_partition, &
      & realspace_cutoff, 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_Allreduce, MPI_Comm_rank, MPI_Comm_size, &
      & MPI_Finalize, MPI_Init
   implicit none

   type(structure_type) :: mol
   type(error_type), allocatable :: error
   integer, allocatable :: num(:)
   real(r8), allocatable :: xyz(:, :), gradient(:, :)
   real(r8) :: energy
   type(d4_model) :: disp
   class(damping_param), allocatable :: 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)

   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("pbe0", param, s9=1.0_r8)
   if (.not.allocated(param)) call fatal("No parameters for PBE0 available")
   call new_d4_model(error, disp, mol)
   if (allocated(error)) call fatal(error%message)

   call new_work_partition(error, partition, rank, nranks)
   if (allocated(error)) call fatal(error%message)

   allocate(gradient(3, mol%nat))
   call get_dispersion(mol, disp, param, realspace_cutoff(), energy, gradient, &
      & partition=partition)

   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 "(a, f13.10, a)", "PBE0-D4 dispersion energy: ", 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 <stdbool.h>
#include <stdio.h>

#include <mpi.h>

#include "dftd4.h"

int main(int argc, char **argv)
{
  dftd4_error error = dftd4_new_error();
  dftd4_structure mol = NULL;
  dftd4_model d4 = NULL;
  dftd4_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] = {
     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 = dftd4_new_structure(error, nat, num, xyz, NULL, NULL, NULL);
  if (dftd4_check_error(error)) goto handle_error;

  param = dftd4_load_rational_damping(error, "pbe0", true);
  if (dftd4_check_error(error)) goto handle_error;

  d4 = dftd4_new_d4_model(error, mol);
  if (dftd4_check_error(error)) goto handle_error;

  dftd4_set_model_work_partition(error, d4, rank, nranks);
  if (dftd4_check_error(error)) goto handle_error;

  double energy;
  double gradient[15];
  dftd4_get_dispersion(error, mol, d4, param, &energy, gradient, NULL);
  if (dftd4_check_error(error)) goto handle_error;

  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("PBE0-D4 dispersion energy: %13.10lf Hartree\n", energy);
  }

  dftd4_delete(error);
  dftd4_delete(mol);
  dftd4_delete(d4);
  dftd4_delete(param);
  MPI_Finalize();
  return 0;

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

  dftd4_delete(error);
  dftd4_delete(mol);
  dftd4_delete(d4);
  dftd4_delete(param);
  MPI_Abort(MPI_COMM_WORLD, 1);
  return 1;
}
}

The examples can be compiled using mpifort or mpicc together with the flags reported by pkg-config dftd4 mctc-lib --cflags --libs. The result is independent of the number of ranks up to the summation order.

Running with one part is identical to omitting the partition. In Fortran this is also available as serial_work_partition; in C and Python the model can be reset with part=0 and nparts=1.

The numerical Hessian uses the same partition. The pairwise decomposition and the model properties do not.