Bug fixes for gpu execution (#330)

* Bug fixes for gpu execution
 - init_gpu should be called very early before mechanism registration
   to make sure all memory is allocated on specific gpu in case of multi-gpu
   system (e.g. bbcore_read & acc copyin)
 - nrnran123_normal should be marked as acc routine seq
 - minor debug message improvement
* fix MPI & GPU initialisation

Co-authored-by: Ioannis <[email protected]>

Author: pramodk

coreneuron/apps/main1.cpp

@@ -138,7 +138,6 @@ void call_prcellstate_for_prcellgid(int prcellgid, int compute_gpu, int is_init)
 void nrn_init_and_load_data(int argc,
                             char* argv[],
                             bool is_mapping_needed = false,
-                            bool nrnmpi_under_nrncontrol = true,
                             bool run_setup_cleanup = true) {
 #if defined(NRN_FEEXCEPT)
     nrn_feenableexcept();
@@ -147,11 +146,6 @@ void nrn_init_and_load_data(int argc,
     /// profiler like tau/vtune : do not measure from begining
     Instrumentor::stop_profile();
 
-// mpi initialisation
-#if NRNMPI
-    nrnmpi_init(nrnmpi_under_nrncontrol ? 1 : 0, &argc, &argv);
-#endif
-
     // memory footprint after mpi initialisation
     report_mem_usage("After MPI_Init");
 
@@ -427,18 +421,26 @@ using namespace coreneuron;
 
 
 extern "C" void mk_mech_init(int argc, char** argv) {
-#if NRNMPI
-    nrnmpi_init(1, &argc, &argv);
-#endif
     // read command line parameters and parameter config files
-
     try {
         corenrn_param.parse(argc, argv);
     }
     catch (...) {
         nrn_abort(1);
     }
 
+#if NRNMPI
+    if (corenrn_param.mpi_enable) {
+        nrnmpi_init(1, &argc, &argv);
+    }
+#endif
+
+#ifdef _OPENACC
+    if (corenrn_param.gpu) {
+        init_gpu();
+    }
+#endif
+
     if (!corenrn_param.writeParametersFilepath.empty()) {
         std::ofstream out(corenrn_param.writeParametersFilepath, std::ios::trunc);
         out << corenrn_param.app.config_to_str(false, false);

coreneuron/gpu/nrn_acc_manager.cpp

@@ -25,13 +25,18 @@ extern InterleaveInfo* interleave_info;
 void copy_ivoc_vect_to_device(IvocVect*& iv, IvocVect*& div);
 void nrn_ion_global_map_copyto_device();
 void nrn_VecPlay_copyto_device(NrnThread* nt, void** d_vecplay);
-void init_gpu(int nthreads, NrnThread* threads);
 
 /* note: threads here are corresponding to global nrn_threads array */
 void setup_nrnthreads_on_device(NrnThread* threads, int nthreads) {
 #ifdef _OPENACC
+    // initialize NrnThreads for gpu execution
+    // empty thread or only artificial cells should be on cpu
+    for (int i = 0; i < nthreads; i++) {
+        NrnThread* nt = threads + i;
+        nt->compute_gpu = (nt->end > 0) ? 1 : 0;
+        nt->_dt = dt;
+    }
 
-    init_gpu(nthreads, threads);
     nrn_ion_global_map_copyto_device();
 
 #ifdef UNIFIED_MEMORY
@@ -946,7 +951,7 @@ void nrn_ion_global_map_copyto_device() {
     }
 }
 
-void init_gpu(int nthreads, NrnThread* threads) {
+void init_gpu() {
     // choose nvidia GPU by default
     acc_device_t device_type = acc_device_nvidia;
 
@@ -969,13 +974,6 @@ void init_gpu(int nthreads, NrnThread* threads) {
     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++) {
-        // empty thread or only artificial cells should be on cpu
-        NrnThread* nt = threads + i;
-        nt->compute_gpu = (nt->end > 0) ? 1 : 0;
-        nt->_dt = dt;
-    }
 }
 
 void nrn_VecPlay_copyto_device(NrnThread* nt, void** d_vecplay) {

coreneuron/gpu/nrn_acc_manager.hpp

@@ -20,6 +20,7 @@ void update_matrix_to_gpu(NrnThread* _nt);
 void update_net_receive_buffer(NrnThread* _nt);
 void realloc_net_receive_buffer(NrnThread* nt, Memb_list* ml);
 void update_net_send_buffer_on_host(NrnThread* nt, NetSendBuffer_t* nsb);
+void init_gpu();
 
 }  // namespace coreneuron
 #endif  // _nrn_device_manager_

coreneuron/io/nrn_setup.hpp

@@ -53,7 +53,6 @@ static void setup_ThreadData(NrnThread& nt);
 extern void nrn_init_and_load_data(int argc,
                                    char** argv,
                                    bool is_mapping_needed = false,
-                                   bool nrnmpi_under_nrncontrol = true,
                                    bool run_setup_cleanup = true);
 extern void nrn_setup_cleanup();
 

coreneuron/network/partrans_setup.cpp

@@ -256,7 +256,9 @@ void nrn_partrans::gap_thread_setup(NrnThread& nt) {
 }
 
 void nrn_partrans::gap_indices_permute(NrnThread& nt) {
-    printf("nrn_partrans::gap_indices_permute\n");
+    if (nrnmpi_myid == 0) {
+        printf("nrn_partrans::gap_indices_permute\n");
+    }
     nrn_partrans::TransferThreadData& ttd = transfer_thread_data_[nt.id];
     // sources
     if (ttd.nsrc > 0 && nt._permute) {

coreneuron/utils/randoms/nrnran123.h

@@ -151,6 +151,9 @@ extern DEVICE double nrnran123_negexp(nrnran123_State*); /* mean 1.0 */
 /* nrnran123_negexp min value is 2.3283064e-10, max is 22.18071 */
 
 /* missing declaration in coreneuron */
+#if !defined(DISABLE_OPENACC)
+#pragma acc routine seq
+#endif
 extern DEVICE double nrnran123_normal(nrnran123_State*);
 
 extern DEVICE double nrnran123_gauss(nrnran123_State*); /* mean 0.0, std 1.0 */