Explicitly initialize HIP/CUDA before MPI

Co-authored-by: Mikael Simberg <mikael.simberg@iki.fi>
cp2k · Nov 21, 2023 · 183bfa0 · 183bfa0
1 parent 65a7a05
commit 183bfa0
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 4 deletions.
diff --git a/src/f77_interface.F b/src/f77_interface.F
@@ -112,6 +112,7 @@ MODULE f77_interface
    USE nnp_environment_types,           ONLY: nnp_type
    USE offload_api,                     ONLY: offload_get_chosen_device,&
                                               offload_get_device_count,&
+                                              offload_init,&
                                               offload_set_chosen_device
    USE periodic_table,                  ONLY: init_periodic_table
    USE pw_gpu,                          ONLY: pw_gpu_finalize,&
@@ -210,7 +211,7 @@ SUBROUTINE init_cp2k(init_mpi, ierr)
       LOGICAL, INTENT(in)                                :: init_mpi
       INTEGER, INTENT(out)                               :: ierr
 
-      INTEGER                                            :: unit_nr
+      INTEGER                                            :: offload_device_count, unit_nr
       TYPE(cp_logger_type), POINTER                      :: logger
 
       IF (.NOT. module_initialized) THEN
@@ -226,6 +227,10 @@ SUBROUTINE init_cp2k(init_mpi, ierr)
          ! get runtime information
          CALL get_runtime_info()
 
+         ! Intialize CUDA/HIP before MPI
+         ! Needed for HIP on ALPS & LUMI
+         CALL offload_init()
+
          ! re-create the para_env and log with correct (reordered) ranks
          ALLOCATE (default_para_env)
          IF (init_mpi) THEN
@@ -262,15 +267,17 @@ SUBROUTINE init_cp2k(init_mpi, ierr)
          !   *** init the bibliography ***
          CALL add_all_references()
 
+         offload_device_count = offload_get_device_count()
+
          ! Select active offload device when available.
-         IF (offload_get_device_count() > 0) THEN
-            CALL offload_set_chosen_device(MOD(default_para_env%mepos, offload_get_device_count()))
+         IF (offload_device_count > 0) THEN
+            CALL offload_set_chosen_device(MOD(default_para_env%mepos, offload_device_count))
          END IF
 
          ! Initialize the DBCSR configuration
          ! Attach the time handler hooks to DBCSR
 #if defined __DBCSR_ACC
-         IF (offload_get_device_count() > 0) THEN
+         IF (offload_device_count > 0) THEN
             CALL dbcsr_init_lib(default_para_env%get_handle(), timeset_hook, timestop_hook, &
                                 cp_abort_hook, cp_warn_hook, io_unit=unit_nr, &
                                 accdrv_active_device_id=offload_get_chosen_device())

diff --git a/src/offload/offload_api.F b/src/offload/offload_api.F
@@ -28,6 +28,7 @@ MODULE offload_api
 
    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'offload_api'
 
+   PUBLIC :: offload_init
    PUBLIC :: offload_get_device_count
    PUBLIC :: offload_set_chosen_device, offload_get_chosen_device, offload_activate_chosen_device
    PUBLIC :: offload_timeset, offload_timestop, offload_mem_info
@@ -86,6 +87,22 @@ END FUNCTION offload_free_pinned_mem_c
       res = offload_free_pinned_mem_c(buffer)
    END FUNCTION offload_free_pinned_mem
 
+! **************************************************************************************************
+!> \brief Initialize runtime.
+!> \return ...
+!> \author Rocco Meli
+! **************************************************************************************************
+   SUBROUTINE offload_init()
+      INTERFACE
+         SUBROUTINE offload_init_c() &
+            BIND(C, name="offload_init")
+         END SUBROUTINE offload_init_c
+      END INTERFACE
+
+      CALL offload_init_c()
+
+   END SUBROUTINE offload_init
+
 ! **************************************************************************************************
 !> \brief Returns the number of available devices.
 !> \return ...

diff --git a/src/offload/offload_library.c b/src/offload/offload_library.c
@@ -14,6 +14,12 @@
 #include "offload_library.h"
 #include "offload_runtime.h"
 
+#if defined(__OFFLOAD_CUDA)
+#include <cuda.h>
+#elif defined(__OFFLOAD_HIP)
+#include <hip/hip_runtime_api.h>
+#endif
+
 #if defined(__OFFLOAD_PROFILING)
 #if defined(__OFFLOAD_CUDA)
 #include <nvToolsExt.h>
@@ -39,6 +45,23 @@ const uint32_t colormap[] = {0xFFFFFF00,  // Yellow
                              0xFF0000FF,  // Blue
                              0xFF000080}; // Navy
 
+/*******************************************************************************
+ * \brief Initialize runtime.
+ * \author Rocco Meli
+ ******************************************************************************/
+void offload_init(void) {
+#if defined(__OFFLOAD_CUDA)
+  CUresult error = cuInit(0);
+  if (error != CUDA_SUCCESS) {
+    fprintf(stderr, "ERROR: %s %d %s %d\n", "cuInit failed with error: ", error,
+            __FILE__, __LINE__);
+    abort();
+  }
+#elif defined(__OFFLOAD_HIP)
+  OFFLOAD_CHECK(hipInit(0));
+#endif
+}
+
 /*******************************************************************************
  * \brief Returns the number of available devices.
  * \author Ole Schuett

diff --git a/src/offload/offload_library.h b/src/offload/offload_library.h
@@ -12,6 +12,12 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
+/*******************************************************************************
+ * \brief Initalize runtime.
+ * \author Rocco Meli
+ ******************************************************************************/
+void offload_init(void);
+
 /*******************************************************************************
  * \brief Returns the number of available devices.
  * \author Ole Schuett