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 "./datatypes.hpp"
#include "./lazy.hpp"
#include "./utils.hpp"
 
#include <mpi.h>
 
#include <type_traits>
#include <utility>
#include <vector>
 
namespace mpi {
 
  namespace detail {
 
    // Type trait to check if a type is a std::vector.
    template <typename T> inline constexpr bool is_std_vector = false;
 
    // Spezialization of is_std_vector for std::vector<T>.
    template <typename T> inline constexpr bool is_std_vector<std::vector<T>> = true;
 
    // Convert an object of type V to an object of type T.
    template <typename T, typename V> T convert(V v) {
      if constexpr (is_std_vector<T>) {
        T res;
        res.reserve(v.size());
        for (auto &x : v) res.emplace_back(convert<typename T::value_type>(std::move(x)));
        return res;
      } else
        return T{std::move(v)};
    }
 
  } // namespace detail
 
  template <typename T> [[gnu::always_inline]] void broadcast(T &&x, communicator c = {}, int root = 0) {
    static_assert(not std::is_const_v<T>, "mpi::broadcast cannot be called on const objects");
    if (has_env) mpi_broadcast(std::forward<T>(x), c, root);
  }
 
  template <typename T>
  [[gnu::always_inline]] inline decltype(auto) reduce(T &&x, communicator c = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    // return type of mpi_reduce
    using r_t = decltype(mpi_reduce(std::forward<T>(x), c, root, all, op));
    if constexpr (is_mpi_lazy<r_t>) {
      return mpi_reduce(std::forward<T>(x), c, root, all, op);
    } else {
      if (has_env)
        return mpi_reduce(std::forward<T>(x), c, root, all, op);
      else
        return detail::convert<r_t>(std::forward<T>(x));
    }
  }
 
  template <typename T>
  [[gnu::always_inline]] inline void reduce_in_place(T &&x, communicator c = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    static_assert(not std::is_const_v<T>, "In-place mpi functions cannot be called on const objects");
    if (has_env) mpi_reduce_in_place(std::forward<T>(x), c, root, all, op);
  }
 
  template <typename T> [[gnu::always_inline]] inline decltype(auto) scatter(T &&x, mpi::communicator c = {}, int root = 0) {
    // return type of mpi_scatter
    using r_t = decltype(mpi_scatter(std::forward<T>(x), c, root));
    if constexpr (is_mpi_lazy<r_t>) {
      return mpi_scatter(std::forward<T>(x), c, root);
    } else {
      if (has_env)
        return mpi_scatter(std::forward<T>(x), c, root);
      else
        return detail::convert<r_t>(std::forward<T>(x));
    }
  }
 
  template <typename T> [[gnu::always_inline]] inline decltype(auto) gather(T &&x, mpi::communicator c = {}, int root = 0, bool all = false) {
    // return type of mpi_gather
    using r_t = decltype(mpi_gather(std::forward<T>(x), c, root, all));
    if constexpr (is_mpi_lazy<r_t>) {
      return mpi_gather(std::forward<T>(x), c, root, all);
    } else {
      if (has_env)
        return mpi_gather(std::forward<T>(x), c, root, all);
      else
        return detail::convert<r_t>(std::forward<T>(x));
    }
  }
 
  template <typename T> [[gnu::always_inline]] inline decltype(auto) all_reduce(T &&x, communicator c = {}, MPI_Op op = MPI_SUM) {
    return reduce(std::forward<T>(x), c, 0, true, op);
  }
 
  template <typename T> [[gnu::always_inline]] inline void all_reduce_in_place(T &&x, communicator c = {}, MPI_Op op = MPI_SUM) {
    reduce_in_place(std::forward<T>(x), c, 0, true, op);
  }
 
  template <typename T> [[gnu::always_inline]] inline decltype(auto) all_gather(T &&x, communicator c = {}) {
    return gather(std::forward<T>(x), c, 0, true);
  }
 
  template <typename T>
    requires(has_mpi_type<T>)
  void mpi_broadcast(T &x, communicator c = {}, int root = 0) {
    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) {
    T b;
    auto d = mpi_type<T>::get();
    if (!all)
      // old MPI implementations may require a non-const send buffer
      check_mpi_call(MPI_Reduce(const_cast<T *>(&x), &b, 1, d, op, root, c.get()), "MPI_Reduce"); // NOLINT
    else
      check_mpi_call(MPI_Allreduce(const_cast<T *>(&x), &b, 1, d, op, c.get()), "MPI_Allreduce"); // NOLINT
    return b;
  }
 
  template <typename T>
    requires(has_mpi_type<T>)
  void mpi_reduce_in_place(T &x, communicator c = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    if (!all)
      check_mpi_call(MPI_Reduce((c.rank() == root ? MPI_IN_PLACE : &x), &x, 1, mpi_type<T>::get(), op, root, c.get()), "MPI_Reduce");
    else
      check_mpi_call(MPI_Allreduce(MPI_IN_PLACE, &x, 1, mpi_type<T>::get(), op, c.get()), "MPI_Allreduce");
  }
 
  template <typename T> bool all_equal(T const &x, communicator c = {}) {
    if (!has_env) 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;
  }
 
} // namespace mpi