generic_communication.hpp

Go to the documentation of this file.

// Copyright (c) 2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Alexander Hampel, Olivier Parcollet, Nils Wentzell

 
#pragma once
 
#include "./communicator.hpp"
#include "./datatypes.hpp"
#include "./macros.hpp"
#include "./utils.hpp"
 
#include <mpi.h>
 
#include <algorithm>
#include <concepts>
#include <ranges>
#include <type_traits>
#include <vector>
 
namespace mpi {

  template <typename R>
  concept MPICompatibleRange = std::ranges::contiguous_range<R> && std::ranges::sized_range<R> && has_mpi_type<std::ranges::range_value_t<R>>;


  template <typename T> [[gnu::always_inline]] void broadcast(T &&x, communicator c = {}, int root = 0) { // NOLINT (forwarding is not needed)
    mpi_broadcast(x, c, root);
  }

  template <typename T>
  [[gnu::always_inline]] decltype(auto) reduce(T &&x, communicator c = {}, int root = 0, bool all = false, // NOLINT (forwarding is not needed)
                                               MPI_Op op = MPI_SUM) {
    if constexpr (requires { mpi_reduce(x, c, root, all, op); }) {
      return mpi_reduce(x, c, root, all, op);
    } else {
      std::remove_cvref_t<T> res;
      reduce_into(x, res, c, root, all, op);
      return res;
    }
  }

  template <typename T>
  [[gnu::always_inline]] void reduce_in_place(T &&x, communicator c = {}, int root = 0, bool all = false, // NOLINT (forwarding is not needed)
                                              MPI_Op op = MPI_SUM) {
    mpi_reduce_into(x, x, c, root, all, op);
  }

  template <typename T1, typename T2>
  [[gnu::always_inline]] void reduce_into(T1 &&x_in, T2 &&x_out, communicator c = {}, int root = 0, // NOLINT (forwarding is not needed)
                                          bool all = false, MPI_Op op = MPI_SUM) {
    mpi_reduce_into(x_in, x_out, c, root, all, op);
  }

  template <typename T>
  [[gnu::always_inline]] decltype(auto) scatter(T &&x, mpi::communicator c = {}, int root = 0) { // NOLINT (forwarding is not needed)
    if constexpr (requires { mpi_scatter(x, c, root); }) {
      return mpi_scatter(x, c, root);
    } else {
      std::remove_cvref_t<T> res;
      scatter_into(x, res, c, root);
      return res;
    }
  }

  template <typename T1, typename T2>
  [[gnu::always_inline]] void scatter_into(T1 &&x_in, T2 &&x_out, communicator c = {}, int root = 0) { // NOLINT (forwarding is not needed)
    mpi_scatter_into(x_in, x_out, c, root);
  }

  template <typename T>
  [[gnu::always_inline]] decltype(auto) gather(T &&x, communicator c = {}, int root = 0, bool all = false) { // NOLINT (forwarding is not needed)
    if constexpr (requires { mpi_gather(x, c, root, all); }) {
      return mpi_gather(x, c, root, all);
    } else {
      std::remove_cvref_t<T> res;
      gather_into(x, res, c, root, all);
      return res;
    }
  }

  template <typename T1, typename T2>
  [[gnu::always_inline]] void gather_into(T1 &&x_in, T2 &&x_out, communicator c = {}, int root = 0, // NOLINT (forwarding is not needed)
                                          bool all = false) {
    mpi_gather_into(x_in, x_out, c, root, all);
  }

  template <typename T>
  [[gnu::always_inline]] decltype(auto) all_reduce(T &&x, communicator c = {}, MPI_Op op = MPI_SUM) { // NOLINT (forwarding is not needed)
    return reduce(x, c, 0, true, op);
  }

  template <typename T>
  [[gnu::always_inline]] void all_reduce_in_place(T &&x, communicator c = {}, MPI_Op op = MPI_SUM) { // NOLINT (forwarding is not needed)
    reduce_in_place(x, c, 0, true, op);
  }

  template <typename T1, typename T2>
  [[gnu::always_inline]] void all_reduce_into(T1 &&x_in, T2 &&x_out, communicator c = {}, MPI_Op op = MPI_SUM) { // NOLINT (forwarding is not needed)
    return reduce_into(x_in, x_out, c, 0, true, op);
  }

  template <typename T> [[gnu::always_inline]] decltype(auto) all_gather(T &&x, communicator c = {}) { // NOLINT (forwarding is not needed)
    return gather(x, c, 0, true);
  }

  template <typename T1, typename T2>
  [[gnu::always_inline]] void all_gather_into(T1 &&x_in, T2 &&x_out, communicator c = {}) { // NOLINT (forwarding is not needed)
    return gather_into(x_in, x_out, c, 0, true);
  }

  template <typename T> bool all_equal(T const &x, communicator c = {}) {
    if (!has_env || c.size() < 2) return true;
    auto min_obj = all_reduce(x, c, MPI_MIN);
    auto max_obj = all_reduce(x, c, MPI_MAX);
    return min_obj == max_obj;
  }

  template <typename T>
    requires(has_mpi_type<T>)
  void mpi_broadcast(T &x, communicator c = {}, int root = 0) {
    // in case there is no active MPI environment or if the communicator size is < 2, do nothing
    if (!has_env || c.size() < 2) return;
 
    // make the MPI C library call
    check_mpi_call(MPI_Bcast(&x, 1, mpi_type<T>::get(), root, c.get()), "MPI_Bcast");
  }

  template <typename T>
    requires(has_mpi_type<T>)
  T mpi_reduce(T const &x, communicator c = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    // in case there is no active MPI environment or if the communicator size is < 2, return the input object
    if (!has_env || c.size() < 2) return x;
 
    // make the MPI C library call with a default constructed output object
    T res;
    if (all) {
      check_mpi_call(MPI_Allreduce(&x, &res, 1, mpi_type<T>::get(), op, c.get()), "MPI_Allreduce");
    } else {
      check_mpi_call(MPI_Reduce(&x, &res, 1, mpi_type<T>::get(), op, root, c.get()), "MPI_Reduce");
    }
    return res;
  }

  template <typename T>
    requires(has_mpi_type<T>)
  void mpi_reduce_into(T const &x_in, T &x_out, communicator c = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    // check if the reduction is in place
    auto in_ptr         = static_cast<void const *>(&x_in);
    auto out_ptr        = static_cast<void *>(&x_out);
    bool const in_place = (in_ptr == out_ptr);
    if (all) {
      EXPECTS_WITH_MESSAGE(all_equal(static_cast<int>(in_place), c),
                           "Either zero or all receiving processes have to choose the in place option in mpi_reduce_into");
    }
 
    // in case there is no active MPI environment or if the communicator size is < 2, do nothing (in place) or copy
    if (!has_env || c.size() < 2) {
      if (!in_place) x_out = x_in;
      return;
    }
 
    // make the MPI C library call
    if (in_place && (c.rank() == root || all)) in_ptr = MPI_IN_PLACE;
    if (all) {
      check_mpi_call(MPI_Allreduce(in_ptr, out_ptr, 1, mpi_type<T>::get(), op, c.get()), "MPI_Allreduce");
    } else {
      check_mpi_call(MPI_Reduce(in_ptr, out_ptr, 1, mpi_type<T>::get(), op, root, c.get()), "MPI_Reduce");
    }
  }

  template <typename T>
    requires(has_mpi_type<T>)
  std::vector<T> mpi_gather(T const &x, communicator c = {}, int root = 0, bool all = false) {
    std::vector<T> res(c.rank() == root || all ? c.size() : 0);
    mpi_gather_into(x, res, c, root, all);
    return res;
  }

  template <typename T, MPICompatibleRange R>
    requires(has_mpi_type<T> && std::same_as<T, std::remove_cvref_t<std::ranges::range_value_t<R>>>)
  void mpi_gather_into(T const &x, R &&rg, communicator c = {}, int root = 0, bool all = false) { // NOLINT (ranges need not be forwarded)
    // check the size of the output range
    if (c.rank() == root || all) {
      EXPECTS_WITH_MESSAGE(c.size() == std::ranges::size(rg), "Output range size is not equal the number of ranks in mpi_gather_into");
    }
 
    // in case there is no active MPI environment or if the communicator size is < 2, copy the input into the range
    if (!has_env || c.size() < 2) {
      std::ranges::copy(std::views::single(x), std::ranges::begin(rg));
      return;
    }
 
    // make the MPI C library call
    using value_t = std::ranges::range_value_t<R>;
    if (all) {
      check_mpi_call(MPI_Allgather(&x, 1, mpi_type<T>::get(), std::ranges::data(rg), 1, mpi_type<value_t>::get(), c.get()), "MPI_Allgather");
    } else {
      check_mpi_call(MPI_Gather(&x, 1, mpi_type<T>::get(), std::ranges::data(rg), 1, mpi_type<value_t>::get(), root, c.get()), "MPI_Gather");
    }
  }

 
} // namespace mpi