UCP/PERF: Merge branch 'master' into ucp-perf-warp-fixes

Author: iyastreb

buildlib/az-helpers.sh

@@ -199,14 +199,14 @@ try_load_cuda_env() {
         have_cuda="${cuda_local_dir}"
     else
         # Fallback to env module
-        az_module_load dev/cuda13.0.0 || return 0
+        az_module_load dev/cuda13.0.2 || return 0
         have_cuda=yes
     fi
 
     # Check gdrcopy
     if [ -w "/dev/gdrdrv" ]
     then
-        az_module_load dev/gdrcopy2.5.1_cuda13.0.0 && have_gdrcopy=yes
+        az_module_load dev/gdrcopy2.5.1_cuda13.0.2 && have_gdrcopy=yes
     fi
 }
 

buildlib/tools/common.sh

@@ -4,8 +4,8 @@ WORKSPACE=${WORKSPACE:=$PWD}
 # build in local directory which goes away when docker exits
 ucx_build_dir=$HOME/${BUILD_ID}/build
 ucx_inst=$ucx_build_dir/install
-CUDA_MODULE="dev/cuda13.0.0"
-GDRCOPY_MODULE="dev/gdrcopy2.5.1_cuda13.0.0"
+CUDA_MODULE="dev/cuda13.0.2"
+GDRCOPY_MODULE="dev/gdrcopy2.5.1_cuda13.0.2"
 JDK_MODULE="dev/jdk"
 MVN_MODULE="dev/mvn"
 XPMEM_MODULE="dev/xpmem-90a95a4"

buildlib/tools/perf-common.yml

@@ -13,7 +13,7 @@ steps:
 
       case "${{ parameters.Name }}" in
         "Build-UCX")
-          module="/hpc/local/etc/modulefiles/dev/cuda13.0.0"
+          module="/hpc/local/etc/modulefiles/dev/cuda13.0.2"
           perfxParams=(--skip-run --source-branch $(Build.SourceBranch) --omb-cuda)
           ;;
         "Perf-test-multi-node")

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);

src/uct/ib/mlx5/gdaki/gdaki.cuh

@@ -165,6 +165,23 @@ UCS_F_DEVICE void uct_rc_mlx5_gda_wqe_prepare_put_or_atomic(
     doca_gpu_dev_verbs_store_wqe_seg(dseg_ptr, (uint64_t*)&(dseg));
 }
 
+UCS_F_DEVICE void uct_rc_mlx5_gda_lock(int *lock) {
+    while (atomicCAS(lock, 0, 1) != 0)
+        ;
+#ifdef DOCA_GPUNETIO_VERBS_HAS_FENCE_ACQUIRE_RELEASE_PTX
+    asm volatile("fence.acquire.gpu;");
+#else
+    uint32_t dummy;
+    uint32_t UCS_V_UNUSED val;
+    asm volatile("ld.acquire.gpu.b32 %0, [%1];" : "=r"(val) : "l"(&dummy));
+#endif
+}
+
+UCS_F_DEVICE void uct_rc_mlx5_gda_unlock(int *lock) {
+    cuda::atomic_ref<int, cuda::thread_scope_device> lock_aref(*lock);
+    lock_aref.store(0, cuda::std::memory_order_release);
+}
+
 UCS_F_DEVICE void uct_rc_mlx5_gda_db(uct_rc_gdaki_dev_ep_t *ep,
                                      uint64_t wqe_base, unsigned count,
                                      uint64_t flags)
@@ -184,13 +201,11 @@ UCS_F_DEVICE void uct_rc_mlx5_gda_db(uct_rc_gdaki_dev_ep_t *ep,
         return;
     }
 
-    doca_gpu_dev_verbs_lock<DOCA_GPUNETIO_VERBS_RESOURCE_SHARING_MODE_GPU>(
-            &ep->sq_lock);
+    uct_rc_mlx5_gda_lock(&ep->sq_lock);
     uct_rc_mlx5_gda_ring_db(ep, ep->sq_ready_index);
     uct_rc_mlx5_gda_update_dbr(ep, ep->sq_ready_index);
     uct_rc_mlx5_gda_ring_db(ep, ep->sq_ready_index);
-    doca_gpu_dev_verbs_unlock<DOCA_GPUNETIO_VERBS_RESOURCE_SHARING_MODE_GPU>(
-            &ep->sq_lock);
+    uct_rc_mlx5_gda_unlock(&ep->sq_lock);
 }
 
 UCS_F_DEVICE bool

src/uct/ib/mlx5/gga/gga_mlx5.c

@@ -852,9 +852,11 @@ uct_gga_mlx5_query_tl_devices(uct_md_h md,
         return UCS_ERR_NO_DEVICE;
     }
 
-    ucs_assertv(mlx5_md->super.cap_flags & UCT_MD_FLAG_EXPORTED_MKEY,
-                "md %p: cap_flags=0x%" PRIx64 " do not have EXPORTED_MKEY flag",
-                mlx5_md, mlx5_md->super.cap_flags);
+    if (!(mlx5_md->super.cap_flags & UCT_MD_FLAG_EXPORTED_MKEY)) {
+        ucs_debug("md %p: cap_flags=0x%" PRIx64 " does not have EXPORTED_MKEY "
+                  "flag", mlx5_md, mlx5_md->super.cap_flags);
+        return UCS_ERR_NO_DEVICE;
+    }
 
     ucs_assertv(ucs_test_all_flags(mlx5_md->flags, UCT_GGA_MLX5_MD_CAPS),
                 "md %p: flags=0x%x do not have mandatory capabilities 0x%x",

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")