Use config flag for locality-aware mpi

hypre-space · Aug 24, 2023 · de4702e · de4702e
1 parent b372b31
commit de4702e
Show file tree

Hide file tree

Showing 10 changed files with 566 additions and 4 deletions.
diff --git a/src/parcsr_ls/par_cycle.c b/src/parcsr_ls/par_cycle.c
@@ -286,6 +286,17 @@ hypre_BoomerAMGCycle( void              *amg_vdata,
    hypre_GpuProfilingPushRange(nvtx_name);
    while (Not_Finished)
    {
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+      if (level >= hypre_HandleNodeAwareSwitchoverThreshold(hypre_handle()))
+      {
+         //if (my_id == 0) { printf("LVL %d using node aware, th: %d\n", level, hypre_HandleNodeAwareSwitchoverThreshold(hypre_handle())); }
+         hypre_HandleUsingNodeAwareMPI(hypre_handle()) = 1;
+      } else
+      {
+         //if (my_id == 0) { printf("LVL %d NOT using node aware\n", level); }
+         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;
@@ -787,6 +803,10 @@ HYPRE_Int hypre_RangeFillResponseIJDetermineRecvProcs ( void *p_recv_contact_buf
 HYPRE_Int hypre_FillResponseIJDetermineSendProcs ( void *p_recv_contact_buf, HYPRE_Int contact_size,
                                                    HYPRE_Int contact_proc, void *ro, MPI_Comm comm, void **p_send_response_buf,
                                                    HYPRE_Int *response_message_size );
+void hypre_ParCSRCreateCommGraph( HYPRE_BigInt first_col_diag,
+                                  HYPRE_BigInt *col_map_offd,
+                                  MPI_Comm comm,
+                                  hypre_ParCSRCommPkg *comm_pkg );
 
 /* numbers.c */
 hypre_NumbersNode *hypre_NumbersNewNode ( void );

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
@@ -750,3 +753,58 @@ hypre_FillResponseIJDetermineSendProcs(void       *p_recv_contact_buf,
 
    return hypre_error_flag;
 }
+
+/*--------------------------------------------------------------------
+ * hypre_ParCSRCreateCommGraph
+ *
+ * Create communication topology graph for MPI neighborhood
+ * collectives
+ *--------------------------------------------------------------------*/
+
+#ifdef HYPRE_USING_NODE_AWARE_MPI
+void
+hypre_ParCSRCreateCommGraph(HYPRE_BigInt first_col_diag,
+                            HYPRE_BigInt *col_map_offd,
+                            MPI_Comm comm,
+                            hypre_ParCSRCommPkg *comm_pkg) {
+   HYPRE_Int num_sends = hypre_ParCSRCommPkgNumSends(comm_pkg);
+   HYPRE_Int num_recvs = hypre_ParCSRCommPkgNumRecvs(comm_pkg);
+   HYPRE_Int *send_map_starts = hypre_ParCSRCommPkgSendMapStarts(comm_pkg);
+   HYPRE_Int *recv_vec_starts = hypre_ParCSRCommPkgRecvVecStarts(comm_pkg);
+   HYPRE_Int *send_map_elmts =  hypre_ParCSRCommPkgSendMapElmts(comm_pkg);
+
+   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_offd[j];
+      }
+   }
+}
+#endif