reduce.hpp

Go to the documentation of this file.

// Copyright (c) 2020-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, Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "./utils.hpp"
#include "../basic_functions.hpp"
#include "../concepts.hpp"
#include "../declarations.hpp"
#include "../exceptions.hpp"
#include "../macros.hpp"
#include "../map.hpp"
#include "../traits.hpp"
 
#include <mpi.h>
#include <mpi/mpi.hpp>
 
#include <array>
#include <cmath>
#include <cstddef>
#include <span>
#include <type_traits>
#include <utility>
 
namespace nda::detail {
 
  // Helper function that (all)reduces arrays/views in-place.
  template <typename A>
    requires(is_regular_or_view_v<A>)
  void mpi_reduce_in_place_impl(A &&a_out, mpi::communicator comm = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) { // NOLINT
    // check the shape of the input arrays/views
    EXPECTS_WITH_MESSAGE(have_mpi_equal_shapes(a_out, comm), "Error in nda::detail::mpi_reduce_in_place_impl: Shapes of arrays/views must be equal")
 
    // do nothing if there is no active MPI environment or if the communicator size is < 2
    if (not mpi::has_env || comm.size() < 2) { return; }
 
    // reduce the arrays/views
    using value_type = typename std::decay_t<A>::value_type;
    if constexpr (not mpi::has_mpi_type<value_type>) {
      // arrays/views of non-MPI types are reduced element-wise
      nda::for_each(a_out.shape(), [&](auto... args) { mpi::reduce_in_place(a_out(args...), comm, root, all, op); });
    } else {
      // for MPI-types we have to perform some checks on the input/ouput arrays/views
      check_layout_mpi_compatible(a_out, "detail::mpi_reduce_in_place_impl");
 
      // reduce the data
      auto a_out_span = std::span{a_out.data(), static_cast<std::size_t>(a_out.size())};
      mpi::reduce_in_place_range(a_out_span, comm, root, all, op);
    }
  }
 
} // namespace nda::detail
 
namespace nda {
 
  template <typename A1, typename A2>
    requires(is_regular_or_view_v<A1> && is_regular_or_view_v<A2>)
  void mpi_reduce_capi(A1 const &a_in, A2 &&a_out, mpi::communicator comm = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) { // NOLINT
    // check the shape of the input arrays/views
    EXPECTS_WITH_MESSAGE(detail::have_mpi_equal_shapes(a_in, comm), "Error in nda::mpi_reduce_capi: Shapes of arrays/views must be equal")
 
    // simply copy if there is no active MPI environment or if the communicator size is < 2
    if (not mpi::has_env || comm.size() < 2) {
      a_out = a_in;
      return;
    }
 
    // reduce the arrays/views
    using value_type = typename std::decay_t<A1>::value_type;
    if constexpr (not mpi::has_mpi_type<value_type>) {
      // arrays/views of non-MPI types are reduced element-wise
      a_out = nda::map([&](auto const &x) { return mpi::reduce(x, comm, root, all, op); })(a_in);
    } else {
      // for MPI-types we have to perform some checks on the input and ouput arrays/views
      detail::check_layout_mpi_compatible(a_in, "mpi_reduce_capi");
      if ((comm.rank() == root) || all) {
        detail::check_layout_mpi_compatible(a_out, "mpi_reduce_capi");
        resize_or_check_if_view(a_out, a_in.shape());
        if (std::abs(a_out.data() - a_in.data()) < a_in.size()) NDA_RUNTIME_ERROR << "Error in nda::mpi_reduce_capi: Overlapping arrays";
      }
 
      // reduce the data
      auto a_out_span = std::span{a_out.data(), static_cast<std::size_t>(a_out.size())};
      auto a_in_span  = std::span{a_in.data(), static_cast<std::size_t>(a_in.size())};
      mpi::reduce_range(a_in_span, a_out_span, comm, root, all, op);
    }
  }
 
  template <typename A>
    requires(is_regular_or_view_v<A>)
  auto lazy_mpi_reduce(A &&a, mpi::communicator comm = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    return mpi::lazy<mpi::tag::reduce, A>{std::forward<A>(a), comm, root, all, op};
  }
 
  template <typename A>
    requires(is_regular_or_view_v<A>)
  void mpi_reduce_in_place(A &&a, mpi::communicator comm = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) { // NOLINT
    detail::mpi_reduce_in_place_impl(a, comm, root, all, op);
  }
 
  template <typename A>
    requires(is_regular_or_view_v<A>)
  auto mpi_reduce(A const &a, mpi::communicator comm = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    using return_t = get_regular_t<A>;
    return_t a_out;
    mpi_reduce_capi(a, a_out, comm, root, all, op);
    return a_out;
  }
 
} // namespace nda
 
template <nda::Array A>
struct mpi::lazy<mpi::tag::reduce, A> {
  using value_type = typename std::decay_t<A>::value_type;
 
  using stored_type = A;
 
  stored_type rhs;
 
  mpi::communicator comm;
 
  const int root{0}; // NOLINT (const is fine here)
 
  const bool all{false}; // NOLINT (const is fine here)
 
  const MPI_Op op{MPI_SUM}; // NOLINT (const is fine here)
 
  [[nodiscard]] auto shape() const {
    if ((comm.rank() == root) || all) return rhs.shape();
    return std::array<long, std::remove_cvref_t<stored_type>::rank>{};
  }
 
  template <nda::Array T>
  void invoke(T &&target) const { // NOLINT (temporary views are allowed here)
    if (target.data() == rhs.data()) {
      nda::detail::mpi_reduce_in_place_impl(target, comm, root, all, op);
    } else {
      nda::mpi_reduce_capi(rhs, target, comm, root, all, op);
    }
  }
};