UCP/DEVICE: Make memh and local_addr optional for counter elements

src/ucp/core/ucp_device.c

@@ -18,6 +18,7 @@
 
 #include "ucp_worker.inl"
 #include "ucp_ep.inl"
+#include "ucp_mm.inl"
 
 
 KHASH_TYPE(ucp_device_handle_allocs, ucp_device_mem_list_handle_h,
@@ -339,6 +340,50 @@ static void ucp_device_mem_list_lane_lookup(
     }
 }
 
+static ucs_status_t
+ucp_device_detect_uct_memh(ucp_context_h context,
+                           const ucp_device_mem_list_elem_t *element,
+                           ucs_memory_type_t mem_type, ucp_md_index_t md_index,
+                           uct_mem_h *uct_memh)
+{
+    void *local_addr;
+    size_t length;
+    ucs_status_t status;
+    ucp_mem_h memh;
+
+    *uct_memh  = UCT_MEM_HANDLE_NULL;
+    local_addr = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD, element,
+                                 local_addr, LOCAL_ADDR, NULL);
+    length     = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD, element,
+                                 length, LENGTH, 0);
+
+    if ((local_addr == NULL) || (length == 0)) {
+        return UCS_OK;
+    }
+
+    status = ucp_memh_get(context, local_addr, length, mem_type,
+                          UCS_BIT(md_index),
+                          UCT_MD_MEM_ACCESS_LOCAL_READ |
+                          UCT_MD_MEM_ACCESS_LOCAL_WRITE,
+                          "device_mem_list", &memh);
+    if (status != UCS_OK) {
+        ucs_debug("failed to detect memh from local_addr %p length %zu: %s",
+                  local_addr, length, ucs_status_string(status));
+        return status;
+    }
+
+    ucs_assertv((memh->md_map & UCS_BIT(md_index)) != 0,
+                "memh=%p md_map=0x%lx md_index=%u", memh, memh->md_map,
+                md_index);
+
+    *uct_memh = memh->uct[md_index];
+    ucs_assert(*uct_memh != UCT_MEM_HANDLE_NULL);
+
+    ucp_memh_put(memh);
+
+    return UCS_OK;
+}
+
 static ucs_status_t ucp_device_mem_list_create_handle(
         ucp_ep_h ep, ucs_sys_device_t local_sys_dev,
         const ucp_device_mem_list_params_t *params,
@@ -365,6 +410,7 @@ static ucs_status_t ucp_device_mem_list_create_handle(
     size_t length;
     void *local_addr;
     uint64_t remote_addr;
+    ucp_mem_h memh;
 
     handle_size += sizeof(*handle.local_addrs) + sizeof(*handle.remote_addrs) +
                    sizeof(*handle.lengths);
@@ -445,6 +491,15 @@ static ucs_status_t ucp_device_mem_list_create_handle(
         ucp_element = &params->elements[i];
         local_addr  = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD,
                                       ucp_element, local_addr, LOCAL_ADDR, NULL);
+
+        if (local_addr == NULL) {
+            memh = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD, ucp_element,
+                                   memh, MEMH, NULL);
+            if (memh != NULL) {
+                local_addr = ucp_memh_address(memh);
+            }
+        }
+
         remote_addr = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD,
                                       ucp_element, remote_addr, REMOTE_ADDR, 0);
         length = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_ELEM_FIELD, ucp_element,
@@ -479,7 +534,8 @@ static ucs_status_t ucp_device_mem_list_create_handle(
                      ucp_element->memh->md_map, local_md_index);
                 ucs_assert(uct_memh != UCT_MEM_HANDLE_NULL);
             } else {
-                uct_memh = UCT_MEM_HANDLE_NULL;
+                ucp_device_detect_uct_memh(ep->worker->context, ucp_element,
+                                           mem_type, local_md_index, &uct_memh);
             }
 
             /* Remote registration */