Scatterer.h

// Copyright (C) 2022 Igor Baratta and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "MPI.h"
#include <algorithm>
#include <mpi.h>
#include <span>
#include <vector>
 
using namespace dolfinx;
 
namespace dolfinx::common
{
class IndexMap;
 
class Scatterer
{
public:
  Scatterer(const IndexMap& map, int bs);
 
  template <typename T>
  void scatter_fwd_begin(const std::span<const T>& send_buffer,
                         const std::span<T>& recv_buffer,
                         MPI_Request& request) const
  {
    // Return early if there are no incoming or outgoing edges
    if (_sizes_local.empty() and _sizes_remote.empty())
      return;
 
    MPI_Ineighbor_alltoallv(
        send_buffer.data(), _sizes_local.data(), _displs_local.data(),
        MPI::mpi_type<T>(), recv_buffer.data(), _sizes_remote.data(),
        _displs_remote.data(), MPI::mpi_type<T>(), _comm0.comm(), &request);
  }
 
  void scatter_fwd_end(MPI_Request& request) const;
 
  template <typename T, typename Functor>
  void scatter_fwd_begin(const std::span<const T>& local_data,
                         std::span<T> local_buffer, std::span<T> remote_buffer,
                         Functor pack_fn, MPI_Request& request) const
  {
    assert(local_buffer.size() == _local_inds.size());
    assert(remote_buffer.size() == _remote_inds.size());
    pack_fn(local_data, _local_inds, local_buffer);
    scatter_fwd_begin(std::span<const T>(local_buffer), remote_buffer, request);
  }
 
  template <typename T, typename Functor>
  void scatter_fwd_end(const std::span<const T>& remote_buffer,
                       std::span<T> remote_data, Functor unpack_fn,
                       MPI_Request& request) const
  {
    assert(remote_buffer.size() == _remote_inds.size());
    assert(remote_data.size() == _remote_inds.size());
    scatter_fwd_end(request);
    unpack_fn(remote_buffer, _remote_inds, remote_data,
              [](T /*a*/, T b) { return b; });
  }
 
  template <typename T>
  void scatter_fwd(const std::span<const T>& local_data,
                   std::span<T> remote_data) const
  {
    MPI_Request request;
    std::vector<T> local_buffer(local_buffer_size(), 0);
    std::vector<T> remote_buffer(remote_buffer_size(), 0);
    auto pack_fn = [](const auto& in, const auto& idx, auto& out)
    {
      for (std::size_t i = 0; i < idx.size(); ++i)
        out[i] = in[idx[i]];
    };
    scatter_fwd_begin(local_data, std::span<T>(local_buffer),
                      std::span<T>(remote_buffer), pack_fn, request);
 
    auto unpack_fn = [](const auto& in, const auto& idx, auto& out, auto op)
    {
      for (std::size_t i = 0; i < idx.size(); ++i)
        out[idx[i]] = op(out[idx[i]], in[i]);
    };
 
    scatter_fwd_end(std::span<const T>(remote_buffer), remote_data, unpack_fn,
                    request);
  }
 
  template <typename T>
  void scatter_rev_begin(const std::span<const T>& send_buffer,
                         const std::span<T>& recv_buffer,
                         MPI_Request& request) const
  {
    // Return early if there are no incoming or outgoing edges
    if (_sizes_local.empty() and _sizes_remote.empty())
      return;
 
    // Send and receive data
    MPI_Ineighbor_alltoallv(
        send_buffer.data(), _sizes_remote.data(), _displs_remote.data(),
        MPI::mpi_type<T>(), recv_buffer.data(), _sizes_local.data(),
        _displs_local.data(), MPI::mpi_type<T>(), _comm1.comm(), &request);
  }
 
  void scatter_rev_end(MPI_Request& request) const;
 
  template <typename T, typename Functor>
  void scatter_rev_begin(const std::span<const T>& remote_data,
                         std::span<T> remote_buffer, std::span<T> local_buffer,
                         Functor pack_fn, MPI_Request& request) const
  {
    assert(local_buffer.size() == _local_inds.size());
    assert(remote_buffer.size() == _remote_inds.size());
    pack_fn(remote_data, _remote_inds, remote_buffer);
    scatter_rev_begin(std::span<const T>(remote_buffer), local_buffer, request);
  }
 
  template <typename T, typename Functor, typename BinaryOp>
  void scatter_rev_end(const std::span<const T>& local_buffer,
                       std::span<T> local_data, Functor unpack_fn, BinaryOp op,
                       MPI_Request& request)
  {
    assert(local_buffer.size() == _local_inds.size());
    if (_local_inds.size() > 0)
      assert(*std::max_element(_local_inds.begin(), _local_inds.end())
             < std::int32_t(local_data.size()));
    scatter_rev_end(request);
    unpack_fn(local_buffer, _local_inds, local_data, op);
  }
 
  template <typename T, typename BinaryOp>
  void scatter_rev(std::span<T> local_data,
                   const std::span<const T>& remote_data, BinaryOp op)
  {
    std::vector<T> local_buffer(local_buffer_size(), 0);
    std::vector<T> remote_buffer(remote_buffer_size(), 0);
    auto pack_fn = [](const auto& in, const auto& idx, auto& out)
    {
      for (std::size_t i = 0; i < idx.size(); ++i)
        out[i] = in[idx[i]];
    };
    auto unpack_fn = [](const auto& in, const auto& idx, auto& out, auto op)
    {
      for (std::size_t i = 0; i < idx.size(); ++i)
        out[idx[i]] = op(out[idx[i]], in[i]);
    };
    MPI_Request request;
    scatter_rev_begin(remote_data, std::span<T>(remote_buffer),
                      std::span<T>(local_buffer), pack_fn, request);
    scatter_rev_end(std::span<const T>(local_buffer), local_data, unpack_fn, op,
                    request);
  }
 
  std::int32_t local_buffer_size() const noexcept;
 
  std::int32_t remote_buffer_size() const noexcept;
 
  const std::vector<std::int32_t>& local_indices() const noexcept;
 
  const std::vector<std::int32_t>& remote_indices() const noexcept;
 
  int bs() const noexcept;
 
private:
  // Block size
  int _bs;
 
  // 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 _comm0;
 
  // 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 _comm1;
 
  // Permutation indices used to pack and unpack ghost data (remote)
  std::vector<std::int32_t> _remote_inds;
 
  // Number of remote indices (ghosts) for each neighbor process
  std::vector<int> _sizes_remote;
 
  // Displacements of remote data for mpi scatter and gather
  std::vector<int> _displs_remote;
 
  // Permutation indices used to pack and unpack local shared data
  // (owned indices that are shared with other processes). Indices are
  // grouped by neighbor process.
  std::vector<std::int32_t> _local_inds;
 
  // Number of local shared indices per neighbor process
  std::vector<int> _sizes_local;
 
  // Displacements of local data for mpi scatter and gather
  std::vector<int> _displs_local;
};
} // namespace dolfinx::common