Add getter for mem list handle num of elements. (#10975)

Author: ofirfarjun7

src/ucp/api/device/ucp_device_types.h

@@ -63,22 +63,22 @@ typedef struct ucp_device_mem_list_handle {
     /**
      * Array of local addresses for the device transfer operations.
      */
-    void     **local_addrs;
+    void            **local_addrs;
 
     /**
      * Array of remote addresses for the device transfer operations.
      */
-    uint64_t *remote_addrs;
-    
+    uint64_t        *remote_addrs;
+
     /**
      * Array of lengths of the local buffers in bytes.
      */
-    size_t *lengths;
+    size_t          *lengths;
 
     /**
      * Array of UCT memory element objects.
      */
-    void *uct_mem_elements;
+    void            *uct_mem_elements;
 
     /**
      * local address, remote address, and length arrays, are allocated contiguously.

src/ucp/api/device/ucp_host.h

@@ -152,7 +152,7 @@ typedef struct ucp_device_mem_list_params {
  *
  * @param [in]  ep        Remote endpoint handle.
  * @param [in]  params    Parameters used to create the handle.
- * @param [out] handle    Created descriptor list handle.
+ * @param [out] handle    Created descriptors list handle.
  *
  * @return Error code as defined by @ref ucs_status_t.
  * @retval UCS_ERR_NOT_CONNECTED if the endpoint is not connected yet.
@@ -164,6 +164,18 @@ ucp_device_mem_list_create(ucp_ep_h ep,
                            ucp_device_mem_list_handle_h *handle);
 
 
+/**
+ * @ingroup UCP_DEVICE
+ * @brief Return the number of elements in the descriptors mem list handle.
+ *
+ * @param [in] handle     Descriptors list handle.
+ *
+ * @return Descriptors mem list length.
+ */
+uint32_t
+ucp_device_get_mem_list_length(const ucp_device_mem_list_handle_h handle);
+
+
 /**
  * @ingroup UCP_DEVICE
  * @brief Release function for a descriptor list handle.

src/ucp/core/ucp_device.c

@@ -21,12 +21,17 @@
 #include "ucp_mm.inl"
 
 
+typedef struct {
+    uct_allocated_memory_t mem;
+    uint32_t               mem_list_length;
+} ucp_device_handle_info_t;
+
 KHASH_TYPE(ucp_device_handle_allocs, ucp_device_mem_list_handle_h,
-           uct_allocated_memory_t);
+           ucp_device_handle_info_t);
 #define ucp_device_handle_hash_key(_handle) \
     kh_int64_hash_func((uintptr_t)(_handle))
 KHASH_IMPL(ucp_device_handle_allocs, ucp_device_mem_list_handle_h,
-           uct_allocated_memory_t, 1, ucp_device_handle_hash_key,
+           ucp_device_handle_info_t, 1, ucp_device_handle_hash_key,
            kh_int64_hash_equal);
 
 /* Hash to track handle allocator, used at release time */
@@ -50,11 +55,16 @@ void ucp_device_cleanup(void)
 }
 
 static ucs_status_t
-ucp_device_mem_handle_hash_insert(uct_allocated_memory_t *mem_handle)
+ucp_device_mem_handle_hash_insert(const uct_allocated_memory_t *mem_handle,
+                                  uint32_t mem_list_length)
 {
     ucs_status_t status;
     khiter_t iter;
     int ret;
+    ucp_device_handle_info_t info;
+
+    info.mem             = *mem_handle;
+    info.mem_list_length = mem_list_length;
 
     ucs_spin_lock(&ucp_device_handle_hash_lock);
     iter = kh_put(ucp_device_handle_allocs, &ucp_device_handle_hash,
@@ -66,7 +76,7 @@ ucp_device_mem_handle_hash_insert(uct_allocated_memory_t *mem_handle)
         ucs_error("handle=%p already found in hash", mem_handle->address);
         status = UCS_ERR_ALREADY_EXISTS;
     } else {
-        kh_value(&ucp_device_handle_hash, iter) = *mem_handle;
+        kh_value(&ucp_device_handle_hash, iter) = info;
         status                                  = UCS_OK;
     }
 
@@ -84,7 +94,7 @@ ucp_device_mem_handle_hash_remove(ucp_device_mem_list_handle_h handle)
     iter = kh_get(ucp_device_handle_allocs, &ucp_device_handle_hash, handle);
     ucs_assertv_always((iter != kh_end(&ucp_device_handle_hash)), "handle=%p",
                        handle);
-    mem = kh_value(&ucp_device_handle_hash, iter);
+    mem = kh_value(&ucp_device_handle_hash, iter).mem;
     kh_del(ucp_device_handle_allocs, &ucp_device_handle_hash, iter);
     ucs_spin_unlock(&ucp_device_handle_hash_lock);
     return mem;
@@ -595,7 +605,7 @@ ucp_device_mem_list_create(ucp_ep_h ep,
     }
 
     /* Track memory allocator for later release */
-    status = ucp_device_mem_handle_hash_insert(&mem);
+    status = ucp_device_mem_handle_hash_insert(&mem, params->num_elements);
     if (status != UCS_OK) {
         uct_mem_free(&mem);
     } else {
@@ -605,6 +615,24 @@ ucp_device_mem_list_create(ucp_ep_h ep,
     return status;
 }
 
+uint32_t
+ucp_device_get_mem_list_length(const ucp_device_mem_list_handle_h handle)
+{
+    khiter_t iter;
+    uint32_t length;
+
+    ucs_assert(handle != NULL);
+
+    ucs_spin_lock(&ucp_device_handle_hash_lock);
+    iter = kh_get(ucp_device_handle_allocs, &ucp_device_handle_hash, handle);
+    ucs_assertv_always((iter != kh_end(&ucp_device_handle_hash)), "handle=%p",
+                       handle);
+    length = kh_value(&ucp_device_handle_hash, iter).mem_list_length;
+    ucs_spin_unlock(&ucp_device_handle_hash_lock);
+
+    return length;
+}
+
 void ucp_device_mem_list_release(ucp_device_mem_list_handle_h handle)
 {
     uct_allocated_memory_t mem = ucp_device_mem_handle_hash_remove(handle);

test/gtest/ucp/test_ucp_device.cc

@@ -334,6 +334,13 @@ UCS_TEST_P(test_ucp_device, create_fail)
     EXPECT_EQ(nullptr, handle);
 }
 
+UCS_TEST_P(test_ucp_device, get_mem_list_length)
+{
+    constexpr unsigned num_elements = 8;
+    mem_list list(sender(), receiver(), 1 * UCS_KBYTE, num_elements);
+    EXPECT_EQ(num_elements, ucp_device_get_mem_list_length(list.handle()));
+}
+
 UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_device, rc_gda, "rc,rc_gda")