Use config flag for locality-aware mpi

hypre-space · Aug 9, 2023 · ef381db · ef381db
1 parent 03b9d3d
commit ef381db
Show file tree

Hide file tree

Showing 8 changed files with 487 additions and 4 deletions.
diff --git a/src/parcsr_ls/par_cycle.c b/src/parcsr_ls/par_cycle.c
@@ -286,6 +286,15 @@ hypre_BoomerAMGCycle( void              *amg_vdata,
    hypre_GpuProfilingPushRange(nvtx_name);
    while (Not_Finished)
    {
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+      if (level >= hypre_HandleNodeAwareSwitchoverThreshold(hypre_handle()))
+      {
+         hypre_HandleUsingNodeAwareMPI(hypre_handle()) = 1;
+      } else
+      {
+         hypre_HandleUsingNodeAwareMPI(hypre_handle()) = 0;
+      }
+#endif
       if (num_levels > 1)
       {
          local_size = hypre_VectorSize(hypre_ParVectorLocalVector(F_array[level]));

diff --git a/src/parcsr_mv/_hypre_parcsr_mv.h b/src/parcsr_mv/_hypre_parcsr_mv.h
@@ -23,6 +23,10 @@ extern "C" {
 #ifndef HYPRE_PAR_CSR_COMMUNICATION_HEADER
 #define HYPRE_PAR_CSR_COMMUNICATION_HEADER
 
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+#include "mpi_advance.h"
+#endif
+
 /*--------------------------------------------------------------------------
  * hypre_ParCSRCommPkg:
  *   Structure containing information for doing communications
@@ -59,13 +63,20 @@ typedef struct
    void                 *recv_data_buffer;
    HYPRE_Int             num_requests;
    hypre_MPI_Request    *requests;
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+   MPIX_Request         *Xrequest;
+#endif
 } hypre_ParCSRCommHandle;
 
 typedef hypre_ParCSRCommHandle hypre_ParCSRPersistentCommHandle;
 
 typedef struct _hypre_ParCSRCommPkg
 {
    MPI_Comm                          comm;
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+   MPIX_Comm                        *neighbor_comm;
+   MPIX_Comm                        *neighborT_comm;
+#endif
    HYPRE_Int                         num_components;
    HYPRE_Int                         num_sends;
    HYPRE_Int                        *send_procs;
@@ -75,6 +86,11 @@ typedef struct _hypre_ParCSRCommPkg
    HYPRE_Int                         num_recvs;
    HYPRE_Int                        *recv_procs;
    HYPRE_Int                        *recv_vec_starts;
+   HYPRE_Int                         use_neighbor;
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+   long                             *global_send_indices;
+   long                             *global_recv_indices;
+#endif
    /* remote communication information */
    hypre_MPI_Datatype               *send_mpi_types;
    hypre_MPI_Datatype               *recv_mpi_types;

diff --git a/src/parcsr_mv/new_commpkg.c b/src/parcsr_mv/new_commpkg.c
@@ -11,6 +11,9 @@
  *-----------------------------------------------------*/
 
 #include "_hypre_parcsr_mv.h"
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+#include <mpi.h>
+#endif
 
 /* some debugging tools*/
 #define mydebug 0
@@ -546,6 +549,43 @@ hypre_ParCSRCommPkgCreateApart
                                     num_sends, send_procs, send_map_starts,
                                     send_map_elmts,
                                     &comm_pkg);
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+   if (comm_pkg->use_neighbor) {
+      int *sendcounts = (int *)malloc(num_sends * sizeof(int));
+      int *recvcounts = (int *)malloc(num_recvs * sizeof(int));
+      for (int i = 0; i < num_sends; i++) {
+         sendcounts[i] = send_map_starts[i+1] - send_map_starts[i];
+      }
+      for (int i = 0; i < num_recvs; i++) {
+         recvcounts[i] = recv_vec_starts[i+1] - recv_vec_starts[i];
+      }
+      MPIX_Dist_graph_create_adjacent( comm, num_recvs, hypre_ParCSRCommPkgRecvProcs(comm_pkg),
+                                       recvcounts,
+                                       num_sends, hypre_ParCSRCommPkgSendProcs(comm_pkg),
+                                       sendcounts,
+                                       MPI_INFO_NULL, 0, &(comm_pkg->neighbor_comm));
+      MPIX_Dist_graph_create_adjacent( comm, num_sends, hypre_ParCSRCommPkgSendProcs(comm_pkg),
+                                       sendcounts,
+                                       num_recvs, hypre_ParCSRCommPkgRecvProcs(comm_pkg),
+                                       recvcounts,
+                                       MPI_INFO_NULL, 0, &(comm_pkg->neighborT_comm));
+
+      HYPRE_Int num_send_elmts = send_map_starts[num_sends];
+      comm_pkg->global_send_indices = hypre_CTAlloc(long, num_send_elmts, HYPRE_MEMORY_HOST);
+      for (int i = 0; i < num_sends; i++) {
+         for (int j = send_map_starts[i]; j < send_map_starts[i+1]; j++) {
+            comm_pkg->global_send_indices[j] = send_map_elmts[j] + first_col_diag;
+         }
+      }
+      HYPRE_Int num_recv_elmts = recv_vec_starts[num_recvs];
+      comm_pkg->global_recv_indices = hypre_CTAlloc(long, num_recv_elmts, HYPRE_MEMORY_HOST);
+      for (int i = 0; i < num_recvs; i++) {
+         for (int j = recv_vec_starts[i]; j < recv_vec_starts[i+1]; j++) {
+            comm_pkg->global_recv_indices[j] = col_map_off_d[j];
+         }
+      }
+   }
+#endif
 
    return hypre_error_flag;
 }