Create parallel communication graph and support NetworkX analysis and visualiation (#3803)

* Work on more detailed IndexMap stats

* more stats

* More stats

* Compile more data

* Remove imbalance

* Remove imbalance from Python

* Provide more detail

* Index map split

* Work on docs and tests

* Add placeholder

* Implement function

* Remove OpenMPI-specific test

* Updates

* Work on edge weights

* Add simple test

* Work on stats

* Work on stats

* Use ranges

* Updates

* Updates

* Simplify

* Type fix

* Enable init

* Improve error message

* Test

* Fix test

* Format

* Enable

* Testing

* Fix

* lint

* Enable

* Re-enable

* Work on stats

* Simplify

* Updates

* Add networkx demo

* mypy fix

* lint

* lint

* Fix type hints

* Fix

* Type fixes

* Update test

* Simplify

* Work on comm plotting

* Improve demo

* Tidy up

* Test update

* Misc fixes

* Doc fixes

* Fixes

* Fixes

* Include fixes

* Add networkx for conda

* Demo name fixes

* Update names

* Fix extension

* Add demo plot options

* Doc fix

* Add bar chart plots

* Remove C++ IndexMap stats code

* Update python/dolfinx/wrappers/graph.cpp

Co-authored-by: Nate <34454754+nate-sime@users.noreply.github.com>

* Update test

* Update cpp/dolfinx/common/IndexMap.h

Co-authored-by: Nate <34454754+nate-sime@users.noreply.github.com>

* Update python/demo/demo_comm-pattern.py

Co-authored-by: Nate <34454754+nate-sime@users.noreply.github.com>

* Update cpp/dolfinx/common/IndexMap.h

Co-authored-by: Nate <34454754+nate-sime@users.noreply.github.com>

* Update node weights

* Doc fixes

* Comment fixes

* More fixes

* Lint

* Use free function

* Simplify

* Fix merge

* Small updates

* Small formatting updates

* Move code to avoid (possible) circular dependencies.

Best of dolfinx::common namespace doesn't depend on other namespaces.

* Tidy

* Add files

* Wrap funtionality

* Work in docs

* Fix imports

* Work in docs

* Update cpp tests

* Doc fix

* Doc fix

---------

Co-authored-by: Chris Richardson <chris@bpi.cam.ac.uk>
Co-authored-by: Nate <34454754+nate-sime@users.noreply.github.com>

Author: 3 people

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

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

cpp/demo/hyperelasticity/main.cpp

@@ -29,6 +29,7 @@
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/cell_types.h>
 #include <dolfinx/nls/NewtonSolver.h>
+#include <numbers>
 #include <petscmat.h>
 #include <petscsys.h>
 #include <petscsystypes.h>
@@ -142,8 +143,7 @@ int main(int argc, char* argv[])
   PetscInitialize(&argc, &argv, nullptr, nullptr);
 
   // Set the logging thread name to show the process rank
-  int mpi_rank;
-  MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
+  int mpi_rank = dolfinx::MPI::rank(MPI_COMM_WORLD);
   std::string fmt = "[%Y-%m-%d %H:%M:%S.%e] [RANK " + std::to_string(mpi_rank)
                     + "] [%l] %v";
   spdlog::set_pattern(fmt);
@@ -194,7 +194,7 @@ int main(int argc, char* argv[])
           constexpr U x2_c = 0.5;
 
           // Large angle of rotation (60 degrees)
-          constexpr U theta = 1.04719755;
+          constexpr U theta = std::numbers::pi / 3;
 
           // New coordinates
           std::vector<U> fdata(3 * x.extent(1), 0.0);

cpp/dolfinx/common/IndexMap.cpp

@@ -11,6 +11,7 @@
 #include <cstdint>
 #include <functional>
 #include <numeric>
+#include <ranges>
 #include <span>
 #include <utility>
 #include <vector>
@@ -182,9 +183,9 @@ compute_submap_indices(const IndexMap& imap,
   std::ranges::for_each(indices,
                         [&is_in_submap](auto i) { is_in_submap[i] = 1; });
 
-  // --- Step 1 ---: Send ghost indices in `indices` to their owners
-  // and receive indices owned by this process that are in `indices`
-  // on other processes
+  // --- Step 1 ---: Send ghost indices in `indices` to their owners and
+  // receive indices owned by this process that are in `indices` on
+  // other processes.
   const auto [send_indices, recv_indices, ghost_buffer_pos, send_sizes,
               recv_sizes, send_disp, recv_disp]
       = communicate_ghosts_to_owners(
@@ -198,7 +199,7 @@ compute_submap_indices(const IndexMap& imap,
   // all indices in `recv_indices` will necessarily be in `indices` on
   // this process, and thus other processes must own them in the submap.
   // If ownership of received index doesn't change, then this process
-  // has the receiving rank as a destination
+  // has the receiving rank as a destination.
   std::vector<int> recv_owners(send_disp.back());
   std::vector<int> submap_dest;
   submap_dest.reserve(1);
@@ -571,7 +572,7 @@ common::compute_owned_indices(std::span<const std::int32_t> indices,
   std::vector<int> send_sizes(src.size(), 0);
   std::vector<int> send_disp(src.size() + 1, 0);
   auto it = ghost_owners.begin();
-  for (std::size_t i = 0; i < src.size(); i++)
+  for (std::size_t i = 0; i < src.size(); ++i)
   {
     int owner = src[i];
     auto begin = std::find(it, ghost_owners.end(), owner);
@@ -850,7 +851,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)
@@ -1112,7 +1112,6 @@ graph::AdjacencyList<int> IndexMap::index_to_dest_ranks(int tag) const
   //    (index, [sharing ranks]). Non-owned indices are ghosted but
   //    not owned by this rank.
   {
-
     // Send data for owned indices back to ghosting ranks (this is
     // necessary to share with ghosting ranks all the ranks that also
     // ghost a ghost index)
@@ -1302,35 +1301,76 @@ 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 = dolfinx::MPI::size(comm_s);
+  int rank = dolfinx::MPI::rank(_comm.comm());
+
+  // 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)};
 }
 //-----------------------------------------------------------------------------

cpp/dolfinx/common/IndexMap.h

@@ -7,10 +7,12 @@
 #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 <tuple>
 #include <utility>
 #include <vector>
 
@@ -255,22 +257,56 @@ 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 a 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
+  /// 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;
 
 private:
   // Range of indices (global) owned by this process
@@ -294,4 +330,5 @@ class IndexMap
   // Set of ranks ghost owned indices
   std::vector<int> _dest;
 };
+
 } // namespace dolfinx::common

cpp/dolfinx/common/MPI.h

@@ -14,7 +14,6 @@
 #include <cassert>
 #include <complex>
 #include <cstdint>
-#include <dolfinx/graph/AdjacencyList.h>
 #include <numeric>
 #include <set>
 #include <span>

cpp/dolfinx/common/Scatterer.h

@@ -23,9 +23,9 @@ namespace dolfinx::common
 /// @brief A Scatterer supports the MPI scattering and gathering of data
 /// that is associated with a common::IndexMap.
 ///
-/// Scatter and gather operations use MPI neighbourhood collectives.
-/// The implementation is designed for sparse communication patterns,
-/// as it typical of patterns based on an IndexMap.
+/// Scatter and gather operations use MPI neighbourhood collectives. The
+/// implementation is designed for sparse communication patterns, as it
+/// typical of patterns based on an IndexMap.
 template <class Allocator = std::allocator<std::int32_t>>
 class Scatterer
 {
@@ -395,7 +395,7 @@ class Scatterer
     if (_sizes_local.empty() and _sizes_remote.empty())
       return;
 
-    // // Send and receive data
+    // Send and receive data
 
     switch (type)
     {
@@ -411,16 +411,17 @@ class Scatterer
     case type::p2p:
     {
       assert(requests.size() == _dest.size() + _src.size());
-      // Start non-blocking send from this process to ghost owners.
+
+      // Start non-blocking send from this process to ghost owners
       for (std::size_t i = 0; i < _dest.size(); i++)
       {
         MPI_Irecv(recv_buffer.data() + _displs_local[i], _sizes_local[i],
                   dolfinx::MPI::mpi_t<T>, _dest[i], MPI_ANY_TAG, _comm0.comm(),
                   &requests[i]);
       }
 
-      // Start non-blocking receive from neighbor process for which an owned
-      // index is a ghost.
+      // Start non-blocking receive from neighbor process for which an
+      // owned index is a ghost
       for (std::size_t i = 0; i < _src.size(); i++)
       {
         MPI_Isend(send_buffer.data() + _displs_remote[i], _sizes_remote[i],
@@ -440,8 +441,8 @@ class Scatterer
   /// The buffers passed to Scatterer::scatter_rev_begin must not be
   /// modified until after the function has been called.
   ///
-  /// @param[in] request The handle used when calling
-  /// Scatterer::scatter_rev_begin
+  /// @param[in] request Handle used when calling
+  /// Scatterer::scatter_rev_begin.
   void scatter_rev_end(std::span<MPI_Request> request) const
   {
     // Return early if there are no incoming or outgoing edges
@@ -568,28 +569,33 @@ class Scatterer
     return _remote_inds.size();
   }
 
-  /// Return a vector of local indices (owned) used to pack/unpack local
-  /// data. These indices are grouped by neighbor process (process for
-  /// which an index is a ghost).
+  /// @brief Return a vector of local indices (owned) used to
+  /// pack/unpack local data.
+  ///
+  /// Indices are grouped by neighbor process (process for which an
+  /// index is a ghost).
   const std::vector<std::int32_t>& local_indices() const noexcept
   {
     return _local_inds;
   }
 
-  /// Return a vector of remote indices (ghosts) used to pack/unpack ghost
-  /// data. These indices are grouped by neighbor process (ghost owners).
+  /// @brief Return a vector of remote indices (ghosts) used to
+  /// pack/unpack ghost data.
+  ///
+  /// These indices are grouped by neighbor process (ghost owners).
   const std::vector<std::int32_t>& remote_indices() const noexcept
   {
     return _remote_inds;
   }
 
   /// @brief The number values (block size) to send per index in the
-  /// common::IndexMap use to create the scatterer
+  /// common::IndexMap use to create the scatterer.
   /// @return The block size
   int bs() const noexcept { return _bs; }
 
-  /// @brief Create a vector of MPI_Requests for a given Scatterer::type
-  /// @return A vector of MPI requests
+  /// @brief Create a vector of MPI_Requests for a given
+  /// Scatterer::type.
+  /// @return Vector of MPI requests.
   std::vector<MPI_Request> create_request_vector(Scatterer::type type
                                                  = type::neighbor)
   {