gather.hpp

Go to the documentation of this file.

// Copyright (c) 2020-2023 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: Olivier Parcollet, Nils Wentzell
 
#pragma once
 
#include "../basic_functions.hpp"
#include "../concepts.hpp"
#include "../exceptions.hpp"
#include "../traits.hpp"
 
#include <mpi/mpi.hpp>
 
#include <type_traits>
#include <utility>
#include <vector>
 
template <nda::Array A>
struct mpi::lazy<mpi::tag::gather, A> {
  using value_type = typename std::decay_t<A>::value_type;
 
  using const_view_type = decltype(std::declval<const A>()());
 
  const_view_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 {
    auto dims = rhs.shape();
    long dim0 = dims[0];
    if (!all) {
      dims[0] = mpi::reduce(dim0, comm, root);
      if (comm.rank() != root) dims[0] = 1;
    } else
      dims[0] = mpi::all_reduce(dim0, comm);
    return dims;
  }
 
  template <nda::Array T>
  void invoke(T &&target) const { // NOLINT (temporary views are allowed here)
    // check if the arrays can be used in the MPI call
    if (not target.is_contiguous()) NDA_RUNTIME_ERROR << "Error in MPI gather for nda::Array: Target array needs to be contiguous";
    static_assert(std::decay_t<A>::layout_t::stride_order_encoded == std::decay_t<T>::layout_t::stride_order_encoded,
                  "Error in MPI gather for nda::Array: Incompatible stride orders");
 
    // special case for non-mpi runs
    if (not mpi::has_env) {
      target = rhs;
      return;
    }
 
    // get target shape and resize or check the target array
    auto dims = shape();
    if (all || (comm.rank() == root)) nda::resize_or_check_if_view(target, dims);
 
    // gather receive counts and memory displacements
    auto recvcounts   = std::vector<int>(comm.size());
    auto displs       = std::vector<int>(comm.size() + 1, 0);
    int sendcount     = rhs.size();
    auto mpi_int_type = mpi::mpi_type<int>::get();
    if (!all)
      MPI_Gather(&sendcount, 1, mpi_int_type, &recvcounts[0], 1, mpi_int_type, root, comm.get());
    else
      MPI_Allgather(&sendcount, 1, mpi_int_type, &recvcounts[0], 1, mpi_int_type, comm.get());
 
    for (int r = 0; r < comm.size(); ++r) displs[r + 1] = recvcounts[r] + displs[r];
 
    // gather the data
    auto mpi_value_type = mpi::mpi_type<value_type>::get();
    if (!all)
      MPI_Gatherv((void *)rhs.data(), sendcount, mpi_value_type, target.data(), &recvcounts[0], &displs[0], mpi_value_type, root, comm.get());
    else
      MPI_Allgatherv((void *)rhs.data(), sendcount, mpi_value_type, target.data(), &recvcounts[0], &displs[0], mpi_value_type, comm.get());
  }
};
 
namespace nda {
 
  template <typename A>
  ArrayInitializer<std::remove_reference_t<A>> auto mpi_gather(A &&a, mpi::communicator comm = {}, int root = 0, bool all = false)
    requires(is_regular_or_view_v<A>)
  {
    if (not a.is_contiguous()) NDA_RUNTIME_ERROR << "Error in MPI gather for nda::Array: Array needs to be contiguous";
    return mpi::lazy<mpi::tag::gather, A>{std::forward<A>(a), comm, root, all};
  }
 
} // namespace nda