diff --git a/src/ucm/rocm/rocmmem.c b/src/ucm/rocm/rocmmem.c index 386f07fdcea..89a9e94e5e9 100644 --- a/src/ucm/rocm/rocmmem.c +++ b/src/ucm/rocm/rocmmem.c @@ -14,9 +14,12 @@ #include #include #include +#include #include #include +#include + #include #include #include @@ -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 index 58d05a29fea..68c3e58e2c4 100644 --- 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 index 1309a2aee76..02135c86103 100644 --- 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));