Scatterer.h

// Copyright (C) 2022-2026 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 "IndexMap.h"
#include "MPI.h"
#include "sort.h"
#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <format>
#include <mpi.h>
#include <numeric>
#include <ranges>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <vector>
 
namespace dolfinx::common
{
template <class Container = std::vector<std::int32_t>>
class Scatterer
{
  static_assert(std::is_integral_v<typename Container::value_type>);
 
  template <class>
  friend class Scatterer;
 
public:
  using container_type = Container;

  explicit Scatterer(const IndexMap& map)
      : _sizes_remote(map.src().size(), 0),
        _displs_remote(map.src().size() + 1), _sizes_local(map.dest().size()),
        _displs_local(map.dest().size() + 1)
  {
    if (dolfinx::MPI::size(map.comm()) == 1)
      return;
 
    int ierr;
    const std::span<const int> src = map.src();
    const std::span<const int> dest = map.dest();
 
    // Check that src and dest ranks are unique and sorted
    assert(std::ranges::is_sorted(src));
    assert(std::ranges::is_sorted(dest));
 
    // Create communicators with directed edges:
    // (0) owner -> ghost,
    // (1) ghost -> owner
    MPI_Comm comm0;
    ierr = MPI_Dist_graph_create_adjacent(
        map.comm(), src.size(), src.data(), MPI_UNWEIGHTED, dest.size(),
        dest.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm0);
    _comm0 = dolfinx::MPI::Comm(comm0, false);
    dolfinx::MPI::check_error(map.comm(), ierr);
 
    MPI_Comm comm1;
    ierr = MPI_Dist_graph_create_adjacent(
        map.comm(), dest.size(), dest.data(), MPI_UNWEIGHTED, src.size(),
        src.data(), MPI_UNWEIGHTED, MPI_INFO_NULL, false, &comm1);
    _comm1 = dolfinx::MPI::Comm(comm1, false);
    dolfinx::MPI::check_error(map.comm(), ierr);
 
    // Build permutation array that sorts ghost indices by owning rank
    std::span owners = map.owners();
    std::vector<std::int32_t> perm(owners.size());
    std::iota(perm.begin(), perm.end(), 0);
    dolfinx::radix_sort(perm, [&owners](auto index) { return owners[index]; });
 
    // Sort (i) ghost indices and (ii) ghost index owners by rank
    // (using perm array)
    std::span ghosts = map.ghosts();
    std::vector<int> owners_sorted(owners.size());
    std::vector<std::int64_t> ghosts_sorted(owners.size());
    std::ranges::transform(perm, owners_sorted.begin(),
                           [&owners](auto idx) { return owners[idx]; });
    std::ranges::transform(perm, ghosts_sorted.begin(),
                           [&ghosts](auto idx) { return ghosts[idx]; });
 
    // For data associated with ghost indices, packed by owning
    // (neighbourhood) rank, compute sizes and displacements. I.e., when
    // sending ghost index data from this rank to the owning ranks,
    // disp[i] is the first entry in the buffer sent to neighbourhood
    // rank i, and disp[i + 1] - disp[i] is the number of values sent to
    // rank i.
    assert(_sizes_remote.size() == src.size());
    assert(_displs_remote.size() == src.size() + 1);
    auto begin = owners_sorted.begin();
    for (std::size_t i = 0; i < src.size(); i++)
    {
      auto upper = std::ranges::upper_bound(begin, owners_sorted.end(), src[i]);
      std::size_t num_ind = std::ranges::distance(begin, upper);
      _displs_remote[i + 1] = _displs_remote[i] + num_ind;
      _sizes_remote[i] = num_ind;
      begin = upper;
    }
 
    // For data associated with owned indices that are ghosted by other
    // ranks, compute the size and displacement arrays. When sending
    // data associated with ghost indices to the owner, these size and
    // displacement arrays are for the receive buffer.
 
    // Compute sizes and displacements of local data (how many local
    // elements to be sent/received grouped by neighbors)
    assert(_sizes_local.size() == dest.size());
    assert(_displs_local.size() == dest.size() + 1);
    _sizes_remote.reserve(1); // ensure data is not a nullptr
    _sizes_local.reserve(1);  // ensure data is not a nullptr
    ierr
        = MPI_Neighbor_alltoall(_sizes_remote.data(), 1, MPI_INT,
                                _sizes_local.data(), 1, MPI_INT, _comm1.comm());
    dolfinx::MPI::check_error(_comm1.comm(), ierr);
 
    std::inclusive_scan(_sizes_local.begin(), _sizes_local.end(),
                        std::next(_displs_local.begin()));
 
    assert(static_cast<int>(ghosts_sorted.size()) == _displs_remote.back());
 
    // Send ghost global indices to owning rank, and receive owned
    // indices that are ghosts on other ranks
    std::vector<std::int64_t> recv_buffer(_displs_local.back(), 0);
    ierr = MPI_Neighbor_alltoallv(
        ghosts_sorted.data(), _sizes_remote.data(), _displs_remote.data(),
        MPI_INT64_T, recv_buffer.data(), _sizes_local.data(),
        _displs_local.data(), MPI_INT64_T, _comm1.comm());
    dolfinx::MPI::check_error(_comm1.comm(), ierr);
 
    const std::array<std::int64_t, 2> range = map.local_range();
#ifndef NDEBUG
    // Check that all received indices are within the owned range
    std::ranges::for_each(recv_buffer, [range](auto idx)
                          { assert(idx >= range[0] and idx < range[1]); });
#endif
 
    // Sizes, displacements and indices are all in blocks. The block
    // size enters only through the MPI datatype used to send them, and
    // through the caller's pack/unpack.
    {
      // Convert the received indices from global to local numbering
      std::vector<typename container_type::value_type> idx(recv_buffer.size());
      std::ranges::transform(recv_buffer, idx.begin(),
                             [offset = range[0]](auto i) ->
                             typename container_type::value_type
                             { return i - offset; });
      _local_inds = std::move(idx);
    }
 
    _remote_inds = container_type(perm.begin(), perm.end());
  }

  template <class U>
  Scatterer(const Scatterer<U>& s)
      : _comm0(s._comm0), _comm1(s._comm1),
        _remote_inds(s._remote_inds.begin(), s._remote_inds.end()),
        _sizes_remote(s._sizes_remote), _displs_remote(s._displs_remote),
        _local_inds(s._local_inds.begin(), s._local_inds.end()),
        _sizes_local(s._sizes_local), _displs_local(s._displs_local)
  {
  }

  Scatterer(const Scatterer& scatterer) = default;

  Scatterer(Scatterer&& scatterer) = default;

  ~Scatterer() = default;
 
  // Copy assignment (deleted). dolfinx::MPI::Comm cannot be copied into
  // an existing object.
  Scatterer& operator=(const Scatterer& scatterer) = delete;

  Scatterer& operator=(Scatterer&& scatterer) = default;

  template <typename T>
  void scatter_fwd_begin_dtype(const T* send_buffer, T* recv_buffer,
                               MPI_Datatype type, MPI_Request& request) const
  {
    if (!has_neighbours())
    {
      request = MPI_REQUEST_NULL;
      return;
    }
 
    int ierr = MPI_Ineighbor_alltoallv(
        send_buffer, _sizes_local.data(), _displs_local.data(), type,
        recv_buffer, _sizes_remote.data(), _displs_remote.data(), type,
        _comm0.comm(), &request);
    dolfinx::MPI::check_error(_comm0.comm(), ierr);
  }

  template <typename T>
  void scatter_fwd_begin(const T* send_buffer, T* recv_buffer, int bs,
                         MPI_Request& request) const
  {
    // Checked here too, to avoid building a datatype that will not be
    // used
    if (!has_neighbours())
    {
      request = MPI_REQUEST_NULL;
      return;
    }
 
    dolfinx::MPI::Datatype<T> type(bs);
    scatter_fwd_begin_dtype(send_buffer, recv_buffer, type.type(), request);
  }

  void scatter_fwd_end(MPI_Request& request) const
  {
    if (!has_neighbours())
      return;
 
    wait(_comm0, request);
  }

  template <typename T>
  void scatter_rev_begin_dtype(const T* send_buffer, T* recv_buffer,
                               MPI_Datatype type, MPI_Request& request) const
  {
    if (!has_neighbours())
    {
      request = MPI_REQUEST_NULL;
      return;
    }
 
    int ierr = MPI_Ineighbor_alltoallv(
        send_buffer, _sizes_remote.data(), _displs_remote.data(), type,
        recv_buffer, _sizes_local.data(), _displs_local.data(), type,
        _comm1.comm(), &request);
    dolfinx::MPI::check_error(_comm1.comm(), ierr);
  }

  template <typename T>
  void scatter_rev_begin(const T* send_buffer, T* recv_buffer, int bs,
                         MPI_Request& request) const
  {
    // Checked here too, to avoid building a datatype that will not be
    // used
    if (!has_neighbours())
    {
      request = MPI_REQUEST_NULL;
      return;
    }
 
    dolfinx::MPI::Datatype<T> type(bs);
    scatter_rev_begin_dtype(send_buffer, recv_buffer, type.type(), request);
  }

  void scatter_rev_end(MPI_Request& request) const
  {
    if (!has_neighbours())
      return;
 
    wait(_comm1, request);
  }

  const container_type& local_indices_block() const noexcept
  {
    return _local_inds;
  }

  const container_type& remote_indices_block() const noexcept
  {
    return _remote_inds;
  }
 
private:
  // False only on a single rank, where _comm0/_comm1 stay MPI_COMM_NULL
  bool has_neighbours() const noexcept
  {
    return _comm0.comm() != MPI_COMM_NULL;
  }
 
  // Complete a non-blocking request, checking errors against `comm`
  static void wait(const dolfinx::MPI::Comm& comm, MPI_Request& request)
  {
    int ierr = MPI_Wait(&request, MPI_STATUS_IGNORE);
    dolfinx::MPI::check_error(comm.comm(), ierr);
  }
 
  // 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{MPI_COMM_NULL};
 
  // 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 indices
  // - out-edges (dest) are to the owning ranks of my ghost indices
  dolfinx::MPI::Comm _comm1{MPI_COMM_NULL};
 
  // Permutation indices used to pack and unpack ghost data (remote)
  container_type _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.
  container_type _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;
};

template <class Container>
std::vector<std::int32_t>
num_sharing_ranks(const IndexMap& map, const Scatterer<Container>& sc,
                  std::span<const std::int32_t> indices, int bs)
{
  if (bs < 1)
  {
    throw std::invalid_argument(
        std::format("Block size must be positive, not {}.", bs));
  }
 
  const std::int32_t size_local = map.size_local();
  const std::int32_t num_blocks = size_local + map.num_ghosts();
 
  // Throw before the scatters, which all ranks must reach
  for (std::int32_t i : indices)
  {
    if (i < 0 or i / bs >= num_blocks)
      throw std::out_of_range("Local index out of range.");
  }
 
  // Initialised to 1 for this rank. The owner holds one
  // local_indices_block entry per ghosting rank, so it counts
  // occurrences rather than accumulating a reverse scatter.
  std::vector<std::int32_t> count(num_blocks, 1);
  const Container& local = sc.local_indices_block();
  const Container& remote = sc.remote_indices_block();
  for (std::int32_t i : local)
    ++count[i];
 
  std::vector<std::int32_t> send(local.size());
  std::ranges::transform(local, send.begin(),
                         [&count](std::int32_t i) { return count[i]; });
  std::vector<std::int32_t> recv(remote.size());
  MPI_Request request = MPI_REQUEST_NULL;
  sc.scatter_fwd_begin(send.data(), recv.data(), 1, request);
  sc.scatter_fwd_end(request);
  for (std::size_t i = 0; i < remote.size(); ++i)
    count[size_local + remote[i]] = recv[i];
 
  std::vector<std::int32_t> n(indices.size());
  std::ranges::transform(indices, n.begin(), [&count, bs](std::int32_t i)
                         { return count[i / bs]; });
  return n;
}
} // namespace dolfinx::common