mpi.hpp

Go to the documentation of this file.

// Copyright (c) 2020-2021 Simons Foundation
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You may obtain a copy of the License at
//     https://www.gnu.org/licenses/gpl-3.0.txt
//
// Authors: Michel Ferrero, Olivier Parcollet, Nils Wentzell

 
#pragma once
 
#include "./block_gf.hpp"
#include "../../utility/exceptions.hpp"
 
#include <mpi/mpi.hpp>
 
#include <cstddef>
#include <numeric>
#include <string>
#include <type_traits>
 
namespace triqs::gfs {
 
  namespace detail {
 
    // Check the shape of block GFs across all processes.
    template <typename G> bool have_mpi_equal_shape(const G &bg, const mpi::communicator &comm) {
      if constexpr (G::arity == 1) {
        return mpi::all_equal(bg.size(), comm);
      } else {
        return mpi::all_equal(bg.size1(), comm) && mpi::all_equal(bg.size2(), comm);
      }
    }
 
    // Hash the block names for comparison.
    template <typename G> std::size_t hash_names(G const &bg) {
      auto bin_op = [](std::size_t acc, auto const &name) { return acc + std::hash<std::string>{}(name); };
      if constexpr (G::arity == 1) {
        return std::accumulate(bg.block_names().begin(), bg.block_names().end(), std::size_t{0}, bin_op);
      } else {
        uint64_t hash = 0;
        for (auto const &vec : bg.block_names()) { hash += std::accumulate(vec.begin(), vec.end(), std::size_t{0}, bin_op); }
        return hash;
      }
    }
 
  } // namespace detail


  template <BlockGf BG> void mpi_broadcast(BG &&bg, mpi::communicator c, int root) { // NOLINT (temporary views are allowed)
    constexpr bool is_view = std::decay_t<BG>::is_view;
 
    // broadcast block names
    if constexpr (!is_view) {
      // for non-view block GFs, we broadcast the block names directly into the block GF
      mpi::broadcast(bg._block_names, c, root);
    } else {
      // for views, we keep the block names in the block GF but check that they are the same as on the root process
      auto names = bg.block_names();
      mpi::broadcast(names, c, root);
      EXPECTS(bg.block_names() == names);
    }
 
    // broadcast data
    mpi::broadcast(bg.data(), c, root);
  }

  template <BlockGf BG1, BlockGf BG2>
  void mpi_reduce_into(BG1 const &bg_in, BG2 &&bg_out, mpi::communicator c, int root, // NOLINT (temporary views are allowed here)
                       bool all, MPI_Op op) {
    constexpr bool is_view = std::decay_t<BG2>::is_view;
 
    // check the shape and block names of the input block GFs
    EXPECTS(mpi::all_equal(detail::hash_names(bg_in), c));
    if (not detail::have_mpi_equal_shape(bg_in, c))
      TRIQS_RUNTIME_ERROR << "Error in triqs::gfs::mpi_reduce_into: Shapes of input block GFs must be equal";
 
    // assign (check) the block names of the output block GF (view) on receiving ranks
    if ((c.rank() == root || all)) {
      if constexpr (is_view) {
        EXPECTS(bg_in.block_names() == bg_out.block_names());
      } else {
        bg_out._block_names = bg_in.block_names();
      }
    }
 
    // reduce the vector (of vectors) of GF objects
    mpi::reduce_into(bg_in.data(), bg_out.data(), c, root, all, op);
  }

  template <BlockGf BG> auto mpi_reduce(BG const &bg, mpi::communicator c = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    auto res = typename BG::regular_type{};
    mpi_reduce_into(bg, res, c, root, all, op);
    return res;
  }

 
} // namespace triqs::gfs