Merge pull request #7981 from edgargabriel/pr/rocm-v1.12.1-fixes

ROCm fixes - v1.12.x
openucx · Feb 24, 2022 · b1af65e · b1af65e
2 parents ce16c7f + e4a2c84
commit b1af65e
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 3 deletions.
diff --git a/src/ucm/rocm/rocmmem.c b/src/ucm/rocm/rocmmem.c
@@ -14,9 +14,12 @@
 #include <ucm/util/reloc.h>
 #include <ucm/util/replace.h>
 #include <ucs/debug/assert.h>
+#include <ucm/util/sys.h>
 #include <ucs/sys/compiler.h>
 #include <ucs/sys/preprocessor.h>
 
+#include <sys/mman.h>
+
 #include <unistd.h>
 #include <pthread.h>
 #include <stdlib.h>
@@ -176,8 +179,36 @@ static ucs_status_t ucm_rocmmem_install(int events)
     return status;
 }
 
+static int ucm_rocm_scan_regions_cb(void *arg, void *addr, size_t length,
+                                    int prot, const char *path)
+{
+    static const char *rocm_path_pattern = "/dev/dri";
+    ucm_event_handler_t *handler = arg;
+    ucm_event_t event;
+
+    if ((prot & (PROT_READ | PROT_WRITE | PROT_EXEC)) &&
+        strncmp(path, rocm_path_pattern, strlen(rocm_path_pattern))) {
+        return 0;
+    }
+    ucm_debug("dispatching initial memtype allocation for %p..%p %s", addr,
+              UCS_PTR_BYTE_OFFSET(addr, length), path);
+
+    event.mem_type.address  = addr;
+    event.mem_type.size     = length;
+    event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */
+
+    ucm_event_enter();
+    handler->cb(UCM_EVENT_MEM_TYPE_ALLOC, &event, handler->arg);
+    ucm_event_leave();
+
+    return 0;
+}
+
 static void ucm_rocmmem_get_existing_alloc(ucm_event_handler_t *handler)
 {
+    if (handler->events & UCM_EVENT_MEM_TYPE_ALLOC) {
+        ucm_parse_proc_self_maps(ucm_rocm_scan_regions_cb, handler);
+    }
 }
 
 static ucm_event_installer_t ucm_rocm_initializer = {

diff --git a/src/uct/rocm/gdr/rocm_gdr_iface.c b/src/uct/rocm/gdr/rocm_gdr_iface.c
@@ -121,7 +121,8 @@ static UCS_CLASS_INIT_FUNC(uct_rocm_gdr_iface_t, uct_md_h md, uct_worker_h worke
                            const uct_iface_params_t *params,
                            const uct_iface_config_t *tl_config)
 {
-    UCS_CLASS_CALL_SUPER_INIT(uct_base_iface_t, &uct_rocm_gdr_iface_ops, NULL,
+    UCS_CLASS_CALL_SUPER_INIT(uct_base_iface_t, &uct_rocm_gdr_iface_ops, 
+		    	      &uct_base_iface_internal_ops,
                               md, worker, params,
                               tl_config UCS_STATS_ARG(params->stats_root)
                               UCS_STATS_ARG(UCT_ROCM_GDR_TL_NAME));

diff --git a/src/uct/rocm/ipc/rocm_ipc_iface.c b/src/uct/rocm/ipc/rocm_ipc_iface.c
@@ -225,8 +225,9 @@ static UCS_CLASS_INIT_FUNC(uct_rocm_ipc_iface_t, uct_md_h md, uct_worker_h worke
 {
     ucs_status_t status;
 
-    UCS_CLASS_CALL_SUPER_INIT(uct_base_iface_t, &uct_rocm_ipc_iface_ops, NULL,
-                              md, worker, params,
+    UCS_CLASS_CALL_SUPER_INIT(uct_base_iface_t, &uct_rocm_ipc_iface_ops, 
+                              &uct_base_iface_internal_ops,
+			      md, worker, params,
                               tl_config UCS_STATS_ARG(params->stats_root)
                               UCS_STATS_ARG(UCT_ROCM_IPC_TL_NAME));