Merge pull request #11613 from PedramAlizadeh/main

Added lazy initialization of rocm component

Author: edgargabriel

opal/mca/accelerator/rocm/accelerator_rocm.h

@@ -34,6 +34,7 @@
 
 
 #include "opal/mca/accelerator/accelerator.h"
+#include "opal/mca/threads/mutex.h"
 
 typedef struct {
     opal_accelerator_base_component_t super;
@@ -60,4 +61,6 @@ OPAL_DECLSPEC extern int opal_accelerator_rocm_verbose;
 OPAL_DECLSPEC extern size_t opal_accelerator_rocm_memcpyH2D_limit;
 OPAL_DECLSPEC extern size_t opal_accelerator_rocm_memcpyD2H_limit;
 
+OPAL_DECLSPEC extern int opal_accelerator_rocm_lazy_init(void);
+
 #endif

opal/mca/accelerator/rocm/accelerator_rocm_component.c

@@ -28,6 +28,10 @@ int opal_accelerator_rocm_verbose = 0;
 size_t opal_accelerator_rocm_memcpyD2H_limit=1024;
 size_t opal_accelerator_rocm_memcpyH2D_limit=1048576;
 
+/* Initialization lock for lazy rocm initialization */
+static opal_mutex_t accelerator_rocm_init_lock;
+static bool accelerator_rocm_init_complete = false;
+
 hipStream_t opal_accelerator_rocm_MemcpyStream = NULL;
 
 /*
@@ -154,8 +158,42 @@ static int accelerator_rocm_component_register(void)
     return OPAL_SUCCESS;
 }
 
+int opal_accelerator_rocm_lazy_init()
+{
+    int err = OPAL_SUCCESS;
+
+    /* Double checked locking to avoid having to
+     * grab locks post lazy-initialization.  */
+    opal_atomic_rmb();
+    if (true == accelerator_rocm_init_complete) {
+        return OPAL_SUCCESS;
+    }
+    OPAL_THREAD_LOCK(&accelerator_rocm_init_lock);
+
+    /* If already initialized, just exit */
+    if (true == accelerator_rocm_init_complete) {
+        goto out;
+    }
+
+    err = hipStreamCreate(&opal_accelerator_rocm_MemcpyStream);
+    if (hipSuccess != err) {
+        opal_output(0, "Could not create hipStream, err=%d %s\n",
+                err, hipGetErrorString(err));
+        goto out;
+    }
+
+    err = OPAL_SUCCESS;
+    opal_atomic_wmb();
+    accelerator_rocm_init_complete = true;
+out:
+    OPAL_THREAD_UNLOCK(&accelerator_rocm_init_lock);
+    return err;
+}
+
 static opal_accelerator_base_module_t* accelerator_rocm_init(void)
 {
+    OBJ_CONSTRUCT(&accelerator_rocm_init_lock, opal_mutex_t);
+    
     hipError_t err;
 
     if (opal_rocm_runtime_initialized) {
@@ -169,13 +207,6 @@ static opal_accelerator_base_module_t* accelerator_rocm_init(void)
         return NULL;
     }
 
-    err = hipStreamCreate(&opal_accelerator_rocm_MemcpyStream);
-    if (hipSuccess != err) {
-        opal_output(0, "Could not create hipStream, err=%d %s\n",
-                err, hipGetErrorString(err));
-        return NULL;
-    }
-
     opal_atomic_mb();
     opal_rocm_runtime_initialized = true;
 
@@ -192,5 +223,6 @@ static void accelerator_rocm_finalize(opal_accelerator_base_module_t* module)
         opal_accelerator_rocm_MemcpyStream = NULL;
     }
 
+    OBJ_DESTRUCT(&accelerator_rocm_init_lock);
     return;
 }

opal/mca/accelerator/rocm/accelerator_rocm_module.c

@@ -87,6 +87,8 @@ static int mca_accelerator_rocm_check_addr (const void *addr, int *dev_id, uint6
             //We might want to set additional flags in a later iteration.
             //*flags |= MCA_ACCELERATOR_FLAGS_HOST_LDSTR;
             //*flags |= MCA_ACCELERATOR_FLAGS_HOST_ATOMICS;
+            /* First access on a device pointer triggers ROCM support lazy initialization. */
+            opal_accelerator_rocm_lazy_init();
             ret = 1;
         } else if (hipMemoryTypeUnified == srcAttr.memoryType) {
             *flags |= MCA_ACCELERATOR_FLAGS_UNIFIED_MEMORY;