Improve GPU device selection in multi-gpu environment

  - compute local rank within shared memory node
  - find our number of gpu nodes and then ping
    devide to specific node
  - add debug message showing #GPUs per node shared by #ranks

Author: pramodk

coreneuron/gpu/nrn_acc_manager.cpp

@@ -11,6 +11,7 @@
 #include "coreneuron/sim/scopmath/newton_struct.h"
 #include "coreneuron/coreneuron.hpp"
 #include "coreneuron/utils/nrnoc_aux.hpp"
+#include "coreneuron/mpi/nrnmpi.h"
 
 #ifdef _OPENACC
 #include <openacc.h>
@@ -946,14 +947,27 @@ void nrn_ion_global_map_copyto_device() {
 }
 
 void init_gpu(int nthreads, NrnThread* threads) {
-    if (nthreads <= 0) {
-        printf("\n Warning: No threads to copy on GPU! ");
-        return;
+    // choose nvidia GPU by default
+    acc_device_t device_type = acc_device_nvidia;
+
+    // check how many gpu devices available
+    int num_devices = acc_get_num_devices(device_type);
+
+    // if no gpu found, can't run on GPU
+    if (num_devices == 0) {
+        nrn_fatal_error("\n ERROR : Enabled GPU execution but couldn't find NVIDIA GPU! \n");
     }
 
-    /** @todo: currently only checking nvidia gpu */
-    if (acc_get_num_devices(acc_device_nvidia) == 0) {
-        printf("\n WARNING: Enabled GPU execution but couldn't find NVIDIA GPU! \n");
+    // get local rank within a node and assign specific gpu gpu for this node.
+    // multiple threads within the node will use same device.
+    int local_rank = nrnmpi_local_rank();
+    int local_size = nrnmpi_local_size();
+
+    int device_num = local_rank % num_devices;
+    acc_set_device_num(device_num, device_type);
+
+    if (nrnmpi_myid == 0) {
+        std::cout << " Info : " << num_devices << " GPUs shared by " << local_size << " ranks per node\n";
     }
 
     for (int i = 0; i < nthreads; i++) {

coreneuron/mpi/nrnmpi.cpp

@@ -223,4 +223,40 @@ int nrnmpi_initialized() {
     return flag;
 }
 
+/**
+ * Return local mpi rank within a shared memory node
+ *
+ * When performing certain operations, we need to know the rank of mpi
+ * process on a given node. This function uses MPI 3 MPI_Comm_split_type
+ * function and MPI_COMM_TYPE_SHARED key to find out the local rank.
+ */
+int nrnmpi_local_rank() {
+    int local_rank = 0;
+#if NRNMPI
+    MPI_Comm local_comm;
+    MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, nrnmpi_myid_world, MPI_INFO_NULL, &local_comm);
+    MPI_Comm_rank(local_comm, &local_rank);
+    MPI_Comm_free(&local_comm);
+#endif
+    return local_rank;
+}
+
+/**
+ * Return number of ranks launched on single shared memory node
+ *
+ * We use MPI 3 MPI_Comm_split_type function and MPI_COMM_TYPE_SHARED key to
+ * determine number of mpi ranks within a shared memory node..
+ */
+int nrnmpi_local_size() {
+    int local_size = 1;
+#if NRNMPI
+    MPI_Comm local_comm;
+    MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, nrnmpi_myid_world, MPI_INFO_NULL, &local_comm);
+    MPI_Comm_size(local_comm, &local_size);
+    MPI_Comm_free(&local_comm);
+#endif
+    return local_size;
+}
+
+
 }  // namespace coreneuron

coreneuron/mpi/nrnmpi.h

@@ -46,6 +46,8 @@ extern int nrnmpi_myid_bbs;       /* rank in nrn_bbs_comm of rank 0 of a subworl
 extern void nrn_abort(int errcode);
 extern void nrn_fatal_error(const char* msg);
 extern double nrn_wtime(void);
+extern int nrnmpi_local_rank();
+extern int nrnmpi_local_size();
 }  // namespace coreneuron
 
 #if defined(NRNMPI)