IndexMap.h

// Copyright (C) 2015-2019 Chris Richardson, Garth N. Wells and Igor Baratta
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <array>
#include <cstdint>
#include <dolfinx/common/MPI.h>
#include <dolfinx/graph/AdjacencyList.h>
#include <map>
#include <memory>
#include <tuple>
#include <utility>
#include <vector>
#include <xtl/xspan.hpp>
 
namespace dolfinx::common
{
// Forward declaration
class IndexMap;
 
std::vector<int32_t>
compute_owned_indices(const xtl::span<const std::int32_t>& indices,
                      const IndexMap& map);
 
std::tuple<std::int64_t, std::vector<std::int32_t>,
           std::vector<std::vector<std::int64_t>>,
           std::vector<std::vector<int>>>
stack_index_maps(
    const std::vector<
        std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
 
 
class IndexMap
{
public:
  enum class Mode
  {
    insert,
    add
  };
 
  enum class Direction
  {
    reverse, // Ghost to owner
    forward, // Owner to ghost
  };
 
  IndexMap(MPI_Comm comm, std::int32_t local_size);
 
  IndexMap(MPI_Comm comm, std::int32_t local_size,
           const xtl::span<const int>& dest_ranks,
           const xtl::span<const std::int64_t>& ghosts,
           const xtl::span<const int>& src_ranks);
 
private:
  template <typename U, typename V, typename W, typename X>
  IndexMap(std::array<std::int64_t, 2> local_range, std::size_t size_global,
           MPI_Comm comm, U&& comm_owner_to_ghost, U&& comm_ghost_to_owner,
           V&& displs_recv_fwd, V&& ghost_pos_recv_fwd, W&& ghosts,
           X&& shared_indices)
      : _local_range(local_range), _size_global(size_global), _comm(comm),
        _comm_owner_to_ghost(std::forward<U>(comm_owner_to_ghost)),
        _comm_ghost_to_owner(std::forward<U>(comm_ghost_to_owner)),
        _displs_recv_fwd(std::forward<V>(displs_recv_fwd)),
        _ghost_pos_recv_fwd(std::forward<V>(ghost_pos_recv_fwd)),
        _ghosts(std::forward<W>(ghosts)),
        _shared_indices(std::forward<X>(shared_indices))
  {
    _sizes_recv_fwd.resize(_displs_recv_fwd.size() - 1, 0);
    std::adjacent_difference(_displs_recv_fwd.cbegin() + 1,
                             _displs_recv_fwd.cend(), _sizes_recv_fwd.begin());
 
    const std::vector<int32_t>& displs_send = _shared_indices->offsets();
    _sizes_send_fwd.resize(_shared_indices->num_nodes(), 0);
    std::adjacent_difference(displs_send.cbegin() + 1, displs_send.cend(),
                             _sizes_send_fwd.begin());
  }
 
public:
  // Copy constructor
  IndexMap(const IndexMap& map) = delete;
 
  IndexMap(IndexMap&& map) = default;
 
  ~IndexMap() = default;
 
  IndexMap& operator=(IndexMap&& map) = default;
 
  // Copy assignment
  IndexMap& operator=(const IndexMap& map) = delete;
 
  std::array<std::int64_t, 2> local_range() const noexcept;
 
  std::int32_t num_ghosts() const noexcept;
 
  std::int32_t size_local() const noexcept;
 
  std::int64_t size_global() const noexcept;
 
  const std::vector<std::int64_t>& ghosts() const noexcept;
 
  MPI_Comm comm() const;
 
  MPI_Comm comm(Direction dir) const;
 
  void local_to_global(const xtl::span<const std::int32_t>& local,
                       const xtl::span<std::int64_t>& global) const;
 
  void global_to_local(const xtl::span<const std::int64_t>& global,
                       const xtl::span<std::int32_t>& local) const;
 
  std::vector<std::int64_t> global_indices() const;
 
  const graph::AdjacencyList<std::int32_t>&
  scatter_fwd_indices() const noexcept;
 
  const std::vector<std::int32_t>& scatter_fwd_ghost_positions() const noexcept;
 
  std::vector<int> ghost_owners() const;
 
  std::map<std::int32_t, std::set<int>> compute_shared_indices() const;
 
  std::pair<IndexMap, std::vector<std::int32_t>>
  create_submap(const xtl::span<const std::int32_t>& indices) const;
 
  template <typename T>
  void scatter_fwd_begin(const xtl::span<const T>& send_buffer,
                         MPI_Datatype& data_type, MPI_Request& request,
                         const xtl::span<T>& recv_buffer) const
  {
    // Send displacement
    const std::vector<int32_t>& displs_send_fwd = _shared_indices->offsets();
 
    // Return early if there are no incoming or outgoing edges
    if (_displs_recv_fwd.size() == 1 and displs_send_fwd.size() == 1)
      return;
 
    // Get block size
    int n;
    MPI_Type_size(data_type, &n);
    n /= sizeof(T);
    if (static_cast<int>(send_buffer.size()) != n * displs_send_fwd.back())
      throw std::runtime_error("Incompatible send buffer size.");
    if (static_cast<int>(recv_buffer.size()) != n * _displs_recv_fwd.back())
      throw std::runtime_error("Incompatible receive buffer size..");
 
    // Start send/receive
    MPI_Ineighbor_alltoallv(send_buffer.data(), _sizes_send_fwd.data(),
                            displs_send_fwd.data(), data_type,
                            recv_buffer.data(), _sizes_recv_fwd.data(),
                            _displs_recv_fwd.data(), data_type,
                            _comm_owner_to_ghost.comm(), &request);
  }
 
  void scatter_fwd_end(MPI_Request& request) const
  {
    // Return early if there are no incoming or outgoing edges
    const std::vector<int32_t>& displs_send_fwd = _shared_indices->offsets();
    if (_displs_recv_fwd.size() == 1 and displs_send_fwd.size() == 1)
      return;
 
    // Wait for communication to complete
    MPI_Wait(&request, MPI_STATUS_IGNORE);
  }
 
  template <typename T>
  void scatter_fwd(const xtl::span<const T>& local_data,
                   xtl::span<T> remote_data, int n) const
  {
    MPI_Datatype data_type;
    if (n == 1)
      data_type = dolfinx::MPI::mpi_type<T>();
    else
    {
      MPI_Type_contiguous(n, dolfinx::MPI::mpi_type<T>(), &data_type);
      MPI_Type_commit(&data_type);
    }
 
    const std::vector<std::int32_t>& indices = _shared_indices->array();
    std::vector<T> send_buffer(n * indices.size());
    for (std::size_t i = 0; i < indices.size(); ++i)
    {
      std::copy_n(std::next(local_data.cbegin(), n * indices[i]), n,
                  std::next(send_buffer.begin(), n * i));
    }
 
    MPI_Request request;
    std::vector<T> buffer_recv(n * _displs_recv_fwd.back());
    scatter_fwd_begin(xtl::span<const T>(send_buffer), data_type, request,
                      xtl::span<T>(buffer_recv));
    scatter_fwd_end(request);
 
    // Copy into ghost area ("remote_data")
    assert(remote_data.size() == n * _ghost_pos_recv_fwd.size());
    for (std::size_t i = 0; i < _ghost_pos_recv_fwd.size(); ++i)
    {
      std::copy_n(std::next(buffer_recv.cbegin(), n * _ghost_pos_recv_fwd[i]),
                  n, std::next(remote_data.begin(), n * i));
    }
 
    if (n != 1)
      MPI_Type_free(&data_type);
  }
 
  template <typename T>
  void scatter_rev_begin(const xtl::span<const T>& send_buffer,
                         MPI_Datatype& data_type, MPI_Request& request,
                         const xtl::span<T>& recv_buffer) const
  {
    // Get displacement vector
    const std::vector<int32_t>& displs_send_fwd = _shared_indices->offsets();
 
    // Return early if there are no incoming or outgoing edges
    if (_displs_recv_fwd.size() == 1 and displs_send_fwd.size() == 1)
      return;
 
    // Get block size
    int n;
    MPI_Type_size(data_type, &n);
    n /= sizeof(T);
    if (static_cast<int>(send_buffer.size()) != n * _ghosts.size())
      throw std::runtime_error("Inconsistent send buffer size.");
    if (static_cast<int>(recv_buffer.size()) != n * displs_send_fwd.back())
      throw std::runtime_error("Inconsistent receive buffer size.");
 
    // Send and receive data
    MPI_Ineighbor_alltoallv(send_buffer.data(), _sizes_recv_fwd.data(),
                            _displs_recv_fwd.data(), data_type,
                            recv_buffer.data(), _sizes_send_fwd.data(),
                            displs_send_fwd.data(), data_type,
                            _comm_ghost_to_owner.comm(), &request);
  }
 
  void scatter_rev_end(MPI_Request& request) const
  {
    // Return early if there are no incoming or outgoing edges
    const std::vector<int32_t>& displs_send_fwd = _shared_indices->offsets();
    if (_displs_recv_fwd.size() == 1 and displs_send_fwd.size() == 1)
      return;
 
    // Wait for communication to complete
    MPI_Wait(&request, MPI_STATUS_IGNORE);
  }
 
  template <typename T>
  void scatter_rev(xtl::span<T> local_data,
                   const xtl::span<const T>& remote_data, int n,
                   IndexMap::Mode op) const
  {
    MPI_Datatype data_type;
    if (n == 1)
      data_type = dolfinx::MPI::mpi_type<T>();
    else
    {
      MPI_Type_contiguous(n, dolfinx::MPI::mpi_type<T>(), &data_type);
      MPI_Type_commit(&data_type);
    }
 
    // Pack send buffer
    std::vector<T> buffer_send;
    buffer_send.resize(n * _displs_recv_fwd.back());
    for (std::size_t i = 0; i < _ghost_pos_recv_fwd.size(); ++i)
    {
      std::copy_n(std::next(remote_data.cbegin(), n * i), n,
                  std::next(buffer_send.begin(), n * _ghost_pos_recv_fwd[i]));
    }
 
    // Exchange data
    MPI_Request request;
    std::vector<T> buffer_recv(n * _shared_indices->array().size());
    scatter_rev_begin(xtl::span<const T>(buffer_send), data_type, request,
                      xtl::span<T>(buffer_recv));
    scatter_rev_end(request);
 
    // Copy or accumulate into "local_data"
    assert(local_data.size() == n * this->size_local());
    const std::vector<std::int32_t>& shared_indices = _shared_indices->array();
    switch (op)
    {
    case Mode::insert:
      for (std::size_t i = 0; i < shared_indices.size(); ++i)
      {
        std::copy_n(std::next(buffer_recv.cbegin(), n * i), n,
                    std::next(local_data.begin(), n * shared_indices[i]));
      }
      break;
    case Mode::add:
      for (std::size_t i = 0; i < shared_indices.size(); ++i)
      {
        for (int j = 0; j < n; ++j)
          local_data[shared_indices[i] * n + j] += buffer_recv[i * n + j];
      }
      break;
    }
 
    if (n != 1)
      MPI_Type_free(&data_type);
  }
 
private:
  // Range of indices (global) owned by this process
  std::array<std::int64_t, 2> _local_range;
 
  // Number indices across communicator
  std::int64_t _size_global;
 
  // MPI communicator (duplicated of 'input' communicator)
  dolfinx::MPI::Comm _comm;
 
  // Communicator where the source ranks own the indices in the callers
  // halo, and the destination ranks 'ghost' indices owned by the
  // caller. I.e.,
  // - in-edges (src) are from ranks that own my ghosts
  // - out-edges (dest) go to ranks that 'ghost' my owned indices
  dolfinx::MPI::Comm _comm_owner_to_ghost;
 
  // Communicator where the source ranks have ghost indices that are
  // owned by the caller, and the destination ranks are the owners of
  // indices in the callers halo region. I.e.,
  // - in-edges (src) are from ranks that 'ghost' my owned indicies
  // - out-edges (dest) are to the owning ranks of my ghost indices
  dolfinx::MPI::Comm _comm_ghost_to_owner;
 
  // MPI sizes and displacements for forward (owner -> ghost) scatter
  // Note: '_displs_send_fwd' can be got from _shared_indices->offsets()
  std::vector<std::int32_t> _sizes_send_fwd, _sizes_recv_fwd, _displs_recv_fwd;
 
  // Position in the recv buffer for a forward scatter for the ith ghost
  // index (_ghost[i]) entry
  std::vector<std::int32_t> _ghost_pos_recv_fwd;
 
  // Local-to-global map for ghost indices
  std::vector<std::int64_t> _ghosts;
 
  // List of owned local indices that are in the ghost (halo) region on
  // other ranks, grouped by rank in the neighbor communicator
  // (destination ranks in forward communicator and source ranks in the
  // reverse communicator), i.e. `_shared_indices.num_nodes() ==
  // size(_comm_owner_to_ghost)`. The array _shared_indices.offsets() is
  // equivalent to 'displs_send_fwd'.
  std::unique_ptr<graph::AdjacencyList<std::int32_t>> _shared_indices;
};
 
} // namespace dolfinx::common