vector.hpp

Go to the documentation of this file.

// Copyright (c) 2019-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 "./communicator.hpp"
#include "./generic_communication.hpp"
#include "./ranges.hpp"
#include "./utils.hpp"
 
#include <mpi.h>
 
#include <type_traits>
#include <utility>
#include <vector>
 
namespace mpi {


  template <typename T> void mpi_broadcast(std::vector<T> &v, communicator c = {}, int root = 0) {
    auto count = v.size();
    broadcast(count, c, root);
    if (c.rank() != root) v.resize(count);
    broadcast_range(v, c, root);
  }

  template <typename T> auto mpi_reduce(std::vector<T> const &v, communicator c = {}, int root = 0, bool all = false, MPI_Op op = MPI_SUM) {
    using value_type = std::remove_cvref_t<decltype(reduce(std::declval<T>()))>;
    std::vector<value_type> res(c.rank() == root || all ? v.size() : 0);
    reduce_range(v, res, c, root, all, op);
    return res;
  }

  template <typename T1, typename T2>
  void mpi_reduce_into(std::vector<T1> const &v_in, std::vector<T2> &v_out, communicator c = {}, int root = 0, bool all = false,
                       MPI_Op op = MPI_SUM) {
    if ((c.rank() == root || all) && v_out.size() != v_in.size()) v_out.resize(v_in.size());
    reduce_range(v_in, v_out, c, root, all, op);
  }

  template <typename T> void mpi_scatter_into(std::vector<T> const &v_in, std::vector<T> &v_out, communicator c = {}, int root = 0) {
    auto scatter_size = static_cast<int>(v_in.size());
    broadcast(scatter_size, c, root);
    auto const recvcount = chunk_length(scatter_size, c.size(), c.rank());
    if (v_out.size() != recvcount) v_out.resize(recvcount);
    scatter_range(v_in, v_out, scatter_size, c, root);
  }

  template <typename T> void mpi_gather_into(std::vector<T> const &v_in, std::vector<T> &v_out, communicator c = {}, int root = 0, bool all = false) {
    auto const gather_size = mpi::all_reduce(v_in.size(), c);
    if ((c.rank() == root || all) && v_out.size() != gather_size) v_out.resize(gather_size);
    gather_range(v_in, v_out, c, root, all);
  }

 
} // namespace mpi