gather.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 "../layout/range.hpp"
#include "../macros.hpp"
#include "../stdutil/array.hpp"
#include "../traits.hpp"
 
#include <mpi/mpi.hpp>
 
#include <cstddef>
#include <functional>
#include <numeric>
#include <span>
#include <type_traits>
#include <utility>
 
namespace nda::detail {
 
  // Helper function to get the shape and total size of the gathered array/view.
  template <typename A>
    requires(is_regular_or_view_v<A> and std::decay_t<A>::is_stride_order_C())
  auto mpi_gather_shape_impl(A const &a, mpi::communicator comm, int root, bool all) {
    auto dims          = a.shape();
    dims[0]            = mpi::all_reduce(dims[0], comm);
    auto gathered_size = std::accumulate(dims.begin(), dims.end(), 1l, std::multiplies<>());
    if (!all && comm.rank() != root) dims = nda::stdutil::make_initialized_array<dims.size()>(0l);
    return std::make_pair(dims, gathered_size);
  }
 
} // namespace nda::detail
 
namespace nda {


  template <typename A1, typename A2>
    requires(is_regular_or_view_v<A1> and std::decay_t<A1>::is_stride_order_C()
             and is_regular_or_view_v<A2> and std::decay_t<A2>::is_stride_order_C())
  void mpi_gather_capi(A1 const &a_in, A2 &&a_out, mpi::communicator comm = {}, int root = 0, bool all = false) { // NOLINT
    // check the shape of the input arrays/views
    EXPECTS_WITH_MESSAGE(detail::have_mpi_equal_shapes(a_in(nda::range(1), nda::ellipsis{}), comm),
                         "Error in nda::mpi_gather_capi: Shapes of arrays/views must be equal save the first one");
 
    // 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;
    }
 
    // check if the input arrays/views can be used in the MPI call
    detail::check_layout_mpi_compatible(a_in, "mpi_gather_capi");
 
    // get output shape, resize or check the output array/view and prepare output span
    auto [dims, gathered_size] = detail::mpi_gather_shape_impl(a_in, comm, root, all);
    auto a_out_span            = std::span{a_out.data(), 0};
    if (all || (comm.rank() == root)) {
      // check if the output array/view can be used in the MPI call
      detail::check_layout_mpi_compatible(a_out, "mpi_gather_capi");
 
      // resize/check the size of the output array/view
      resize_or_check_if_view(a_out, dims);
 
      // prepare the output span
      a_out_span = std::span{a_out.data(), static_cast<std::size_t>(a_out.size())};
    }
 
    // gather the data
    auto a_in_span = std::span{a_in.data(), static_cast<std::size_t>(a_in.size())};
    mpi::gather_range(a_in_span, a_out_span, gathered_size, comm, root, all);
  }

  template <typename A>
    requires(is_regular_or_view_v<A> and std::decay_t<A>::is_stride_order_C())
  auto lazy_mpi_gather(A &&a, mpi::communicator comm = {}, int root = 0, bool all = false) {
    return mpi::lazy<mpi::tag::gather, A>{std::forward<A>(a), comm, root, all};
  }

  template <typename A>
    requires(is_regular_or_view_v<A> and std::decay_t<A>::is_stride_order_C())
  auto mpi_gather(A const &a, mpi::communicator comm = {}, int root = 0, bool all = false) {
    using return_t = get_regular_t<A>;
    return_t a_out;
    mpi_gather_capi(a, a_out, comm, root, all);
    return a_out;
  }

 
} // namespace nda

template <nda::Array A>
struct mpi::lazy<mpi::tag::gather, 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)

  [[nodiscard]] auto shape() const { return nda::detail::mpi_gather_shape_impl(rhs, comm, root, all).first; }

  template <nda::Array T>
    requires(std::decay_t<T>::is_stride_order_C())
  void invoke(T &&target) const { // NOLINT (temporary views are allowed here)
    nda::mpi_gather_capi(rhs, target, comm, root, all);
  }
};