Create parallel communication graph and support NetworkX analysis and visualiation

.github/workflows/oneapi-conda/environment.yml

@@ -20,6 +20,7 @@ dependencies:
   - mpi4py
   - mpich
   - nanobind
+  - networkx
   - ninja
   - numba
   - numpy

cpp/dolfinx/common/IndexMap.cpp

@@ -9,9 +9,12 @@
 #include "sort.h"
 #include <algorithm>
 #include <cstdint>
+#include <format>
 #include <functional>
 #include <numeric>
+#include <ranges>
 #include <span>
+#include <sstream>
 #include <utility>
 #include <vector>
 
@@ -850,7 +853,6 @@ common::create_sub_index_map(const IndexMap& imap,
                    submap_ghost_gidxs, submap_ghost_owners),
           std::move(sub_imap_to_imap)};
 }
-
 //-----------------------------------------------------------------------------
 //-----------------------------------------------------------------------------
 IndexMap::IndexMap(MPI_Comm comm, std::int32_t local_size) : _comm(comm, true)
@@ -1301,35 +1303,192 @@ std::span<const int> IndexMap::src() const noexcept { return _src; }
 //-----------------------------------------------------------------------------
 std::span<const int> IndexMap::dest() const noexcept { return _dest; }
 //-----------------------------------------------------------------------------
-std::array<double, 2> IndexMap::imbalance() const
+std::vector<std::int32_t> IndexMap::weights_src() const
+{
+  std::vector<std::int32_t> weights(_src.size(), 0);
+  for (int r : _owners)
+  {
+    auto it = std::ranges::lower_bound(_src, r);
+    assert(it != _src.end() and *it == r);
+
+    std::size_t pos = std::distance(_src.begin(), it);
+    assert(pos < weights.size());
+    weights[pos] += 1;
+  }
+
+  return weights;
+}
+//-----------------------------------------------------------------------------
+std::vector<std::int32_t> IndexMap::weights_dest() const
+{
+  int ierr = 0;
+  std::vector<std::int32_t> w_src = this->weights_src();
+
+  std::vector<MPI_Request> requests(_dest.size() + _src.size());
+
+  std::vector<std::int32_t> w_dest(_dest.size());
+  for (std::size_t i = 0; i < _dest.size(); ++i)
+  {
+    ierr = MPI_Irecv(w_dest.data() + i, 1, MPI_INT32_T, _dest[i], MPI_ANY_TAG,
+                     _comm.comm(), &requests[i]);
+    dolfinx::MPI::check_error(_comm.comm(), ierr);
+  }
+
+  for (std::size_t i = 0; i < _src.size(); ++i)
+  {
+    ierr = MPI_Isend(w_src.data() + i, 1, MPI_INT32_T, _src[i], 0, _comm.comm(),
+                     &requests[i + _dest.size()]);
+    dolfinx::MPI::check_error(_comm.comm(), ierr);
+  }
+
+  ierr = MPI_Waitall(requests.size(), requests.data(), MPI_STATUS_IGNORE);
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  return w_dest;
+}
+//-----------------------------------------------------------------------------
+std::array<std::vector<int>, 2> IndexMap::rank_type(int split_type) const
 {
-  std::array<double, 2> imbalance{-1., -1.};
-  std::array<std::int32_t, 2> max_count;
-  std::array<std::int32_t, 2> local_sizes
-      = {static_cast<std::int32_t>(_local_range[1] - _local_range[0]),
-         static_cast<std::int32_t>(_ghosts.size())};
-
-  // Find the maximum number of owned indices and the maximum number of ghost
-  // indices across all processes.
-  MPI_Allreduce(local_sizes.data(), max_count.data(), 2, MPI_INT32_T, MPI_MAX,
-                _comm.comm());
-
-  std::int32_t total_num_ghosts = 0;
-  MPI_Allreduce(&local_sizes[1], &total_num_ghosts, 1, MPI_INT32_T, MPI_SUM,
-                _comm.comm());
-
-  // Compute the average number of owned and ghost indices per process.
-  int comm_size = dolfinx::MPI::size(_comm.comm());
-  double avg_owned = static_cast<double>(_size_global) / comm_size;
-  double avg_ghosts = static_cast<double>(total_num_ghosts) / comm_size;
-
-  // Compute the imbalance by dividing the maximum number of indices by the
-  // corresponding average.
-  if (avg_owned > 0)
-    imbalance[0] = max_count[0] / avg_owned;
-  if (avg_ghosts > 0)
-    imbalance[1] = max_count[1] / avg_ghosts;
-
-  return imbalance;
+  int ierr;
+
+  MPI_Comm comm_s;
+  ierr = MPI_Comm_split_type(_comm.comm(), split_type, 0, MPI_INFO_NULL,
+                             &comm_s);
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  int size_s = 0;
+  ierr = MPI_Comm_size(comm_s, &size_s);
+  dolfinx::MPI::check_error(comm_s, ierr);
+
+  int rank = 0;
+  ierr = MPI_Comm_rank(_comm.comm(), &rank);
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  // Note: in most cases, size_s will be much smaller than the size of
+  // _comm
+  std::vector<int> ranks_s(size_s);
+  ierr = MPI_Allgather(&rank, 1, MPI_INT, ranks_s.data(), 1, MPI_INT, comm_s);
+  dolfinx::MPI::check_error(comm_s, ierr);
+
+  std::vector<int> split_dest, split_src;
+  std::ranges::set_intersection(_dest, ranks_s, std::back_inserter(split_dest));
+  assert(std::ranges::is_sorted(split_dest));
+  std::ranges::set_intersection(_src, ranks_s, std::back_inserter(split_src));
+  assert(std::ranges::is_sorted(split_src));
+
+  ierr = MPI_Comm_free(&comm_s);
+  dolfinx::MPI::check_error(comm_s, ierr);
+
+  return {std::move(split_dest), std::move(split_src)};
+}
+//-----------------------------------------------------------------------------
+graph::AdjacencyList<std::tuple<int, std::size_t, std::int8_t>, std::int32_t>
+IndexMap::comm_graph(int root) const
+{
+  int ierr;
+
+  // Graph edge out(dest) weights
+  const std::vector<std::int32_t> w_dest = this->weights_dest();
+
+  // Group ranks by type
+  const auto [local_dest, local_src] = this->rank_type(MPI_COMM_TYPE_SHARED);
+
+  // Get number of edges for each node (rank)
+  int num_edges_local = _dest.size();
+  std::vector<int> num_edges_remote(dolfinx::MPI::size(_comm.comm()));
+  ierr = MPI_Gather(&num_edges_local, 1, MPI_INT, num_edges_remote.data(), 1,
+                    MPI_INT, root, _comm.comm());
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  // Compute displacements
+  std::vector<int> disp(num_edges_remote.size() + 1, 0);
+  std::partial_sum(num_edges_remote.begin(), num_edges_remote.end(),
+                   std::next(disp.begin()));
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  // For each node (rank), get edge indices
+  std::vector<int> edges_remote(disp.back());
+  edges_remote.reserve(1);
+  ierr = MPI_Gatherv(_dest.data(), _dest.size(), MPI_INT, edges_remote.data(),
+                     num_edges_remote.data(), disp.data(), MPI_INT, root,
+                     _comm.comm());
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  // For each edge, get edge weight
+  std::vector<std::int32_t> weights_remote(disp.back());
+  weights_remote.reserve(1);
+  ierr = MPI_Gatherv(w_dest.data(), w_dest.size(), MPI_INT32_T,
+                     weights_remote.data(), num_edges_remote.data(),
+                     disp.data(), MPI_INT32_T, root, _comm.comm());
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  // For node get local size
+  std::int32_t size = this->size_local();
+  std::vector<std::int32_t> sizes_remote(dolfinx::MPI::size(_comm.comm()));
+  ierr = MPI_Gather(&size, 1, MPI_INT32_T, sizes_remote.data(), 1, MPI_INT32_T,
+                    root, _comm.comm());
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  // For each edge, get its local/remote marker
+  std::vector<std::int8_t> markers;
+  for (auto r : _dest)
+  {
+    auto it = std::ranges::lower_bound(local_dest, r);
+    if (it != local_dest.end() and *it == r)
+      markers.push_back(1);
+    else
+      markers.push_back(0);
+  }
+  std::vector<std::int8_t> markers_remote(disp.back());
+  ierr = MPI_Gatherv(markers.data(), markers.size(), MPI_INT8_T,
+                     markers_remote.data(), num_edges_remote.data(),
+                     disp.data(), MPI_INT8_T, root, _comm.comm());
+  dolfinx::MPI::check_error(_comm.comm(), ierr);
+
+  std::vector<std::tuple<int, std::size_t, std::int8_t>> e_data;
+  for (std::size_t i = 0; i < edges_remote.size(); ++i)
+    e_data.emplace_back(edges_remote[i], weights_remote[i], markers_remote[i]);
+  return graph::AdjacencyList(std::move(e_data),
+                              std::vector(disp.begin(), disp.end()),
+                              std::move(sizes_remote));
+}
+//-----------------------------------------------------------------------------
+std::string IndexMap::comm_to_json(
+    const graph::AdjacencyList<std::tuple<int, std::size_t, std::int8_t>,
+                               std::int32_t>& g)
+{
+  const std::vector<std::int32_t>& node_weights = g.node_data().value();
+
+  std::stringstream out;
+  out << std::format("{{\"directed\": true, \"multigraph\": true, \"graph\": "
+                     "[], \"nodes\": [");
+  for (std::int32_t n = 0; n < g.num_nodes(); ++n)
+  {
+    out << std::format("{{\"weight\": {}, \"id\": {}}}", node_weights[n], n);
+    if (n != g.num_nodes() - 1)
+      out << ", ";
+  }
+  out << "], ";
+  out << "\"adjacency\": [";
+  for (std::int32_t n = 0; n < g.num_nodes(); ++n)
+  {
+    out << "[";
+    auto links = g.links(n);
+    for (std::size_t edge = 0; edge < links.size(); ++edge)
+    {
+      auto [e, w, local] = links[edge];
+      out << std::format(
+          "{{\"local\": {}, \"weight\": {}, \"id\": {}, \"key\": 0}}", local, w,
+          e);
+      if (edge != links.size() - 1)
+        out << ", ";
+    }
+    out << "]";
+    if (n != g.num_nodes() - 1)
+      out << ", ";
+  }
+  out << "]}";
+
+  return out.str();
 }
 //-----------------------------------------------------------------------------

cpp/dolfinx/common/IndexMap.h

@@ -7,10 +7,13 @@
 #pragma once
 
 #include "IndexMap.h"
+#include "MPI.h"
 #include <cstdint>
-#include <dolfinx/common/MPI.h>
+#include <dolfinx/graph/AdjacencyList.h>
 #include <memory>
 #include <span>
+#include <string>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -255,22 +258,101 @@ class IndexMap
   /// and sorted.
   std::span<const int> dest() const noexcept;
 
-  /// @brief Returns the imbalance of the current IndexMap.
+  /// @brief Compute the number of ghost indices owned by each rank in
+  /// IndexMap::src.
   ///
-  /// The imbalance is a measure of load balancing across all processes,
-  /// defined as the maximum number of indices on any process divided by
-  /// the average number of indices per process. This function
-  /// calculates the imbalance separately for owned indices and ghost
-  /// indices and returns them as a std::array<double, 2>. If the total
-  /// number of owned or ghost indices is zero, the respective entry in
-  /// the array is set to -1.
+  /// This is a measure of the amount of data:
   ///
-  /// @note This is a collective operation and must be called by all
-  /// processes in the communicator associated with the IndexMap.
+  /// 1. Sent from this rank to other ranks when performing a reverse
+  /// (owner <- ghost) scatter.
   ///
-  /// @return An array containing the imbalance in owned indices (first
-  /// element) and the imbalance in ghost indices (second element).
-  std::array<double, 2> imbalance() const;
+  /// 2. Received by this rank from other ranks when performing a
+  /// forward (owner -> ghost) scatter.
+  ///
+  /// @return A weight vector, where `weight[i]` the the number of
+  /// ghost indices owned by rank IndexMap::src()`[i]`.
+  std::vector<std::int32_t> weights_src() const;
+
+  /// @brief Compute the number of ghost indices owned by each rank in
+  /// IndexMap::dest.
+  ///
+  /// This is a measure of the amount of data:
+  ///
+  /// 1. Sent from this rank to other ranks when performing a forward
+  /// (owner -> ghost) scatter.
+  ///
+  /// 2. Received by this rank from other ranks when performing a
+  /// reverse forward (owner <- ghost) scatter.
+  ///
+  /// @return A weight vector, where `weight[i]` the the number of ghost
+  /// indices owned by rank IndexMap::dest()`[i]`.
+  std::vector<std::int32_t> weights_dest() const;
+
+  /// @brief Destination and source ranks by type, e.g, ranks that are
+  /// destination/source ranks for the caller and are in an common
+  /// shared memory region.
+  ///
+  /// This function is used to group destination and source ranks by
+  /// 'type'. The type is defined by the MPI `split_type`. Split types
+  /// include ranks from a common shared memory region
+  /// (`MPI_COMM_TYPE_SHARED`) or a common NUMA region. Splits types are
+  /// listed at See
+  /// https://docs.open-mpi.org/en/main/man-openmpi/man3/MPI_Comm_split_type.3.html#split-types.
+  ///
+  /// @note Collective operation on comm();
+  ///
+  /// @param[in] split_type MPI split type, as used in the function
+  /// `MPI_Comm_split_type`. See
+  /// https://docs.open-mpi.org/en/main/man-openmpi/man3/MPI_Comm_split_type.3.html#split-types.
+  /// @return (0) Intersection of ranks in `split_type` and in dest(),
+  /// and (1) intersection of ranks in `split_type` and in src().
+  /// Returned ranks are on the comm() communicator.
+  std::array<std::vector<int>, 2> rank_type(int split_type) const;
+
+  /// @brief Compute an directed graph that describes the parallel
+  /// communication patterns.
+  ///
+  /// The graph describes the communication pattern for a 'forward
+  /// scatter', i.e. sending owned data to ranks that ghost the data
+  /// (owner->ghost operation).
+  ///
+  /// Each node in the graph corresponds to an MPI rank. A graph edge is
+  /// a forward (owner->ghost) communication path. The edge weight is
+  /// the number 'values' communicated along the edge. Each edge also
+  /// has a marker that indicates if the edge is sending data to:
+  ///
+  /// 1. A node (rank) that shares memory with the sender (`true`), or
+  ///
+  /// 2. A remote node that does not share memory with the sender
+  ///   (`false`).
+  ///
+  /// The graph data can be visualised using a tool like
+  /// [NetworkX](https://networkx.org/),
+  ///
+  /// @note Collective.
+  ///
+  /// @param[in] root MPI rank on which the build the communication
+  /// graph data.
+  /// @return Adjacency list representing the communication pattern,
+  /// where edges data is (0) the edge, (1) edge weight and (2)
+  /// local/remote memory indicator, and node weights (number of owned
+  /// indices).
+  graph::AdjacencyList<std::tuple<int, std::size_t, std::int8_t>, std::int32_t>
+  comm_graph(int root = 0) const;
+
+  /// @brief Build communication graph data as JSON string.
+  ///
+  /// The data string can be decoded (loaded) to create a Python object
+  /// from which a [NetworkX](https://networkx.org/) graph can be
+  /// constructed.
+  ///
+  /// See ::comm_graph for a description of the data.
+  ///
+  /// @param[in] g Communication graph.
+  /// @return JSON string representing the communication graph.
+  static std::string comm_to_json(
+      const graph::AdjacencyList<std::tuple<int, std::size_t, std::int8_t>,
+                                 std::int32_t>& g);
 
 private:
   // Range of indices (global) owned by this process
@@ -294,4 +376,5 @@ class IndexMap
   // Set of ranks ghost owned indices
   std::vector<int> _dest;
 };
+
 } // namespace dolfinx::common