ranges.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 "./chunk.hpp"
#include "./communicator.hpp"
#include "./datatypes.hpp"
#include "./environment.hpp"
#include "./generic_communication.hpp"
#include "./macros.hpp"
#include "./utils.hpp"
 
#include <itertools/itertools.hpp>
#include <mpi.h>
 
#include <algorithm>
#include <concepts>
#include <limits>
#include <numeric>
#include <ranges>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace mpi {


  template <std::ranges::sized_range R> void broadcast_range(R &&rg, communicator c = {}, int root = 0) { // NOLINT (ranges need not be forwarded)
    // check the size of the range
    auto size = static_cast<long>(std::ranges::size(rg));
    EXPECTS_WITH_MESSAGE(all_equal(size, c), "Range sizes are not equal on all processes in mpi::broadcast_range");
 
    // do nothing if no elements are broadcasted
    if (size <= 0) return;
 
    // call the MPI C library if the ranges are contiguous with MPI compatible value types, otherwise do element-wise
    // broadcasts
    if constexpr (MPICompatibleRange<R>) {
      // 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 (allow the number of elements to larger than INT_MAX)
      constexpr long max_int = std::numeric_limits<int>::max();
      for (long offset = 0; size > 0; offset += max_int, size -= max_int) {
        auto const count = static_cast<int>(std::min(size, max_int));
        check_mpi_call(MPI_Bcast(std::ranges::data(rg) + offset, count, mpi_type<std::ranges::range_value_t<R>>::get(), root, c.get()), "MPI_Bcast");
      }
    } else {
      // otherwise call the generic broadcast for each element separately
      for (auto &x : rg) broadcast(x, c, root);
    }
  }

  template <std::ranges::sized_range R1, std::ranges::sized_range R2>
  void reduce_range(R1 &&in_rg, R2 &&out_rg, communicator c = {}, int root = 0, bool all = false, // NOLINT (ranges need not be forwarded)
                    MPI_Op op = MPI_SUM) {
    // check the size of the input range
    auto size = static_cast<long>(std::ranges::size(in_rg));
    EXPECTS_WITH_MESSAGE(all_equal(size, c), "Input range sizes are not equal on all processes in mpi::reduce_range");
 
    // do nothing if no elements are reduced
    if (size <= 0) return;
 
    // check the size of the output range
    bool const receives = (c.rank() == root || all);
    if (receives) EXPECTS_WITH_MESSAGE(size == std::ranges::size(out_rg), "Input and output range sizes are not equal in mpi::reduce_range");
 
    // call the MPI C library if the ranges are contiguous with MPI compatible value types
    if constexpr (MPICompatibleRange<R1> && MPICompatibleRange<R2>) {
      static_assert(std::same_as<std::remove_cvref_t<std::ranges::range_value_t<R1>>, std::remove_cvref_t<std::ranges::range_value_t<R2>>>,
                    "Value types of input and output ranges not compatible in mpi::reduce_range");
 
      // check if the reduction is in place
      bool const in_place = (static_cast<void const *>(std::ranges::data(in_rg)) == static_cast<void *>(std::ranges::data(out_rg)));
      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_range");
      }
 
      // in case there is no active MPI environment or if the communicator size is < 2, copy to the output range
      if (!has_env || c.size() < 2) {
        std::ranges::copy(std::forward<R1>(in_rg), std::ranges::data(out_rg));
        return;
      }
 
      // make the MPI C library call (allow the number of elements to larger than INT_MAX)
      constexpr long max_int = std::numeric_limits<int>::max();
      for (long offset = 0; size > 0; offset += max_int, size -= max_int) {
        auto in_data  = static_cast<void const *>(std::ranges::data(in_rg) + offset);
        auto out_data = std::ranges::data(out_rg) + offset;
        if (receives and in_place) in_data = MPI_IN_PLACE;
        auto const count = static_cast<int>(std::min(size, max_int));
        if (all) {
          check_mpi_call(MPI_Allreduce(in_data, out_data, count, mpi_type<std::ranges::range_value_t<R1>>::get(), op, c.get()), "MPI_Allreduce");
        } else {
          check_mpi_call(MPI_Reduce(in_data, out_data, count, mpi_type<std::ranges::range_value_t<R1>>::get(), op, root, c.get()), "MPI_Reduce");
        }
      }
    } else {
      // fallback to element-wise reduction if the range is not contiguous with an MPI compatible value type
      if (size <= std::ranges::size(out_rg)) {
        // on ranks where the output range size is large enough, reduce into the output elements
        for (auto &&[x_in, x_out] : itertools::zip(in_rg, out_rg)) reduce_into(x_in, x_out, c, root, all, op);
      } else {
        // on all other ranks, reduce into a dummy output object (needs to be default constructible)
        using out_value_t = std::ranges::range_value_t<R2>;
        if constexpr (std::is_default_constructible_v<out_value_t>) {
          out_value_t out_dummy{};
          for (auto &&x_in : in_rg) reduce_into(x_in, out_dummy, c, root, all, op);
        } else {
          // if it is not default constructible, is there something we can do?
          throw std::runtime_error("Cannot default construct dummy object in mpi::reduce_range");
        }
      }
    }
  }

  template <MPICompatibleRange R1, MPICompatibleRange R2>
    requires(std::same_as<std::remove_cvref_t<std::ranges::range_value_t<R1>>, std::remove_cvref_t<std::ranges::range_value_t<R2>>>)
  void scatter_range(R1 &&in_rg, R2 &&out_rg, long scatter_size, communicator c = {}, int root = 0, // NOLINT (ranges need not be forwarded)
                     long chunk_size = 1) {
    // check the number of elements to be scattered
    EXPECTS_WITH_MESSAGE(all_equal(scatter_size, c), "Number of elements to be scattered is not equal on all processes in mpi::scatter_range");
 
    // do nothing if no elements are scattered
    if (scatter_size == 0) return;
 
    // check the size of the input range on root
    if (c.rank() == root) {
      EXPECTS_WITH_MESSAGE(scatter_size == std::ranges::size(in_rg),
                           "Input range size on root is not equal the number of elements to be scattered in mpi::scatter_range");
    }
 
    // check the size of the output range
    auto const recvcount = static_cast<int>(chunk_length(scatter_size, c.size(), c.rank(), chunk_size));
    EXPECTS_WITH_MESSAGE(recvcount == std::ranges::size(out_rg),
                         "Output range size is not equal the number of elements to be received in mpi::scatter_range");
 
    // in case there is no active MPI environment or if the communicator size is < 2, copy to output range
    if (!has_env || c.size() < 2) {
      std::ranges::copy(std::forward<R1>(in_rg), std::ranges::data(out_rg));
      return;
    }
 
    // prepare arguments for the MPI call
    auto sendcounts = std::vector<int>(c.size());
    auto displs     = std::vector<int>(c.size() + 1, 0);
    for (int i = 0; i < c.size(); ++i) {
      sendcounts[i] = static_cast<int>(chunk_length(scatter_size, c.size(), i, chunk_size));
      displs[i + 1] = sendcounts[i] + displs[i];
    }
 
    // make the MPI C library call
    check_mpi_call(MPI_Scatterv(std::ranges::data(in_rg), sendcounts.data(), displs.data(), mpi_type<std::ranges::range_value_t<R1>>::get(),
                                std::ranges::data(out_rg), recvcount, mpi_type<std::ranges::range_value_t<R2>>::get(), root, c.get()),
                   "MPI_Scatterv");
  }

  template <MPICompatibleRange R1, MPICompatibleRange R2>
    requires(std::same_as<std::remove_cvref_t<std::ranges::range_value_t<R1>>, std::remove_cvref_t<std::ranges::range_value_t<R2>>>)
  void gather_range(R1 &&in_rg, R2 &&out_rg, communicator c = {}, int root = 0, bool all = false) { // NOLINT (ranges need not be forwarded)
    // get the receive counts (sendcount from each process) and the displacements
    auto sendcount  = static_cast<int>(std::ranges::size(in_rg));
    auto recvcounts = all_gather(sendcount, c);
    auto displs     = std::vector<int>(c.size() + 1, 0);
    std::partial_sum(recvcounts.begin(), recvcounts.end(), displs.begin() + 1);
 
    // do nothing if there are no elements to gather
    if (displs.back() == 0) return;
 
    // check the size of the output range on receiving ranks
    if (c.rank() == root || all) {
      EXPECTS_WITH_MESSAGE(displs.back() == std::ranges::size(out_rg),
                           "Output range size is not equal the number of elements to be received in mpi::gather_range");
    }
 
    // in case there is no active MPI environment or if the communicator size is < 2, copy to the output range
    if (!has_env || c.size() < 2) {
      std::ranges::copy(std::forward<R1>(in_rg), std::ranges::data(out_rg));
      return;
    }
 
    // make the MPI C library call
    if (all) {
      check_mpi_call(MPI_Allgatherv(std::ranges::data(in_rg), sendcount, mpi_type<std::ranges::range_value_t<R1>>::get(), std::ranges::data(out_rg),
                                    recvcounts.data(), displs.data(), mpi_type<std::ranges::range_value_t<R2>>::get(), c.get()),
                     "MPI_Allgatherv");
    } else {
      check_mpi_call(MPI_Gatherv(std::ranges::data(in_rg), sendcount, mpi_type<std::ranges::range_value_t<R1>>::get(), std::ranges::data(out_rg),
                                 recvcounts.data(), displs.data(), mpi_type<std::ranges::range_value_t<R2>>::get(), root, c.get()),
                     "MPI_Gatherv");
    }
  }

 
} // namespace mpi