Use config flag for locality-aware mpi

hypre-space · Jul 3, 2023 · 73daa01 · 73daa01
1 parent 53dfbe3
commit 73daa01
Show file tree

Hide file tree

Showing 4 changed files with 293 additions and 2 deletions.
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,10 @@ typedef struct _hypre_ParCSRCommPkg
    HYPRE_Int                         num_recvs;
    HYPRE_Int                        *recv_procs;
    HYPRE_Int                        *recv_vec_starts;
+#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,48 @@ hypre_ParCSRCommPkgCreateApart
                                     num_sends, send_procs, send_map_starts,
                                     send_map_elmts,
                                     &comm_pkg);
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+   int rank;
+   MPI_Comm_rank(comm, &rank);
+   if (rank ==0) {printf("Comm pkg Create\n");}
+   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),
+                                    //MPI_UNWEIGHTED,
+                                    recvcounts,
+                                    num_sends, hypre_ParCSRCommPkgSendProcs(comm_pkg),
+                                    sendcounts,
+                                    //MPI_UNWEIGHTED,
+                                    MPI_INFO_NULL, 0, &(comm_pkg->neighbor_comm));
+   MPIX_Dist_graph_create_adjacent( comm, num_sends, hypre_ParCSRCommPkgSendProcs(comm_pkg),
+                                    //MPI_UNWEIGHTED,
+                                    sendcounts,
+                                    num_recvs, hypre_ParCSRCommPkgRecvProcs(comm_pkg),
+                                    //MPI_UNWEIGHTED,
+                                    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;
 }