reduce.hpp

Go to the documentation of this file.

// Copyright (c) 2020--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#pragma once
 
#include "./utils.hpp"
#include "../basic_functions.hpp"
#include "../declarations.hpp"
#include "../macros.hpp"
#include "../map.hpp"
#include "../traits.hpp"
 
#include <mpi.h>
#include <mpi/mpi.hpp>
 
#include <cstddef>
#include <span>
#include <type_traits>
#include <utility>
 
namespace nda {


  template <typename A1, typename A2>
    requires(is_regular_or_view_v<A1> && is_regular_or_view_v<A2>)
  void mpi_reduce_into(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_into: Shapes of arrays/views must be equal");
 
    // resize or check the output array/view on receiving ranks
    bool const receives = (all || (comm.rank() == root));
    if (receives) resize_or_check_if_view(a_out, a_in.shape());
 
    // call mpi::reduce_range with a span if the input and output arrays/views are contiguous with positive strides
    // (on non-receiving ranks, the output should always be contiguous since it is not used)
    if (a_in.is_contiguous() and a_in.has_positive_strides()) {
      auto a_in_span = std::span{a_in.data(), static_cast<std::size_t>(a_in.size())};
      if ((a_out.is_contiguous() and a_out.has_positive_strides()) || !receives) {
        auto a_out_span = std::span{a_out.data(), static_cast<std::size_t>(a_out.size())};
        mpi::reduce_range(a_in_span, a_out_span, comm, root, all, op);
      } else {
        mpi::reduce_range(a_in_span, a_out, comm, root, all, op);
      }
    } else {
      if ((a_out.is_contiguous() and a_out.has_positive_strides()) || !receives) {
        auto a_out_span = std::span{a_out.data(), static_cast<std::size_t>(a_out.size())};
        mpi::reduce_range(a_in, a_out_span, comm, root, all, op);
      } else {
        mpi::reduce_range(a_in, a_out, 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 value_t  = std::remove_cvref_t<decltype(mpi::reduce(std::declval<get_value_t<A>>()))>;
    using return_t = basic_array<value_t, get_rank<A>, typename A::layout_policy_t::contiguous_t, get_algebra<A>, heap<>>;
    return_t a_out;
    mpi_reduce_into(a, a_out, comm, root, all, op);
    return a_out;
  }

 
} // namespace nda