UCT/API: Introduce device API (#10829)

Author: ofirfarjun7

buildlib/tools/coverity.sh

@@ -102,7 +102,7 @@ run_coverity() {
 	if [ "${ucx_build_type}" == "devel" ]; then
 		cov-manage-emit --dir $cov_build --tu-pattern "file('.*/test/gtest/common/googletest/*')" delete || :
 	fi
-	cov-analyze --jobs $parallel_jobs $COV_OPT --security --concurrency --dir $cov_build
+	cov-analyze --jobs $parallel_jobs $COV_OPT --disable PARSE_ERROR --security --concurrency --dir $cov_build
 	nerrors=$(cov-format-errors --dir $cov_build | awk '/Processing [0-9]+ errors?/ { print $2 }')
 
 	if [ $nerrors -gt 0 ]; then

contrib/check_inst_headers.sh

@@ -28,7 +28,7 @@ do
 	fi
 
 	# devices files should be ignored for now
-	if test "$hfile" != "${hfile#ucp/api/device/}"
+	if test "$hfile" != "${hfile#uc[pt]/api/device/}"
 	then
 		echo "SKIPPED $hfile (device compiler)"
 		continue

src/tools/info/tl_info.c

@@ -138,6 +138,7 @@ static void print_iface_info(uct_worker_h worker, uct_md_h md,
     };
     uct_iface_config_t *iface_config;
     uct_iface_attr_t iface_attr;
+    uct_iface_attr_v2_t iface_attr_v2;
     char max_eps_str[32];
     ucs_status_t status;
     uct_iface_h iface;
@@ -343,6 +344,15 @@ static void print_iface_info(uct_worker_h worker, uct_md_h md,
         printf("#       error handling:%s\n", buf);
     }
 
+    iface_attr_v2.field_mask = UCT_IFACE_ATTR_FIELD_DEVICE_MEM_ELEMENT_SIZE;
+    status                   = uct_iface_query_v2(iface, &iface_attr_v2);
+    if (status != UCS_OK) {
+        printf("#   < failed to query interface >\n");
+    } else {
+        printf("#   device mem_element: %zu bytes\n",
+               iface_attr_v2.device_mem_element_size);
+    }
+
     uct_iface_close(iface);
     printf("#\n");
 }

src/ucp/api/device/ucp_device_impl.h

@@ -7,12 +7,25 @@
 #ifndef UCP_DEVICE_IMPL_H
 #define UCP_DEVICE_IMPL_H
 
+#include "ucp_host.h"
 #include "ucp_device_types.h"
 
+#include <uct/api/device/uct_device_impl.h>
 #include <ucs/sys/compiler_def.h>
 #include <ucs/type/status.h>
 #include <stdint.h>
 
+/**
+ * @ingroup UCP_DEVICE
+ * @brief GPU request descriptor of a given batch
+ *
+ * This request tracks a batch of memory operations in progress. It can be used
+ * with @ref ucp_device_progress_req to detect request completion.
+ */
+typedef struct ucp_device_request {
+    uct_device_completion_t comp;
+} ucp_device_request_t;
+
 
 /**
  * @ingroup UCP_DEVICE
@@ -226,6 +239,19 @@ ucp_device_put_multi_partial(ucp_device_mem_list_handle_h mem_list,
 }
 
 
+/**
+ * @ingroup UCP_DEVICE
+ *
+ * @brief Initialize a device request.
+ *
+ * @param [out] req  Device request to initialize.
+ */
+UCS_F_DEVICE void ucp_device_request_init(ucp_device_request_t *req)
+{
+    uct_device_completion_init(&req->comp);
+}
+
+
 /**
  * @ingroup UCP_DEVICE
  * @brief Progress a device request containing a batch of operations.
@@ -246,6 +272,11 @@ template <ucp_device_level_t level = UCP_DEVICE_LEVEL_THREAD>
 UCS_F_DEVICE ucs_status_t
 ucp_device_progress_req(ucp_device_request_t *req)
 {
+    if (ucs_likely(req->comp.count == 0)) {
+        return req->comp.status;
+    }
+
+    /* TODO call uct progress function */
     return UCS_ERR_NOT_IMPLEMENTED;
 }
 

src/ucp/api/device/ucp_device_types.h

@@ -7,12 +7,11 @@
 #ifndef UCP_DEVICE_TYPES_H
 #define UCP_DEVICE_TYPES_H
 
-#include <uct/api/uct.h>
+#include <uct/api/device/uct_device_types.h>
+#include <uct/api/uct_def.h>
 
 
-typedef struct ucp_mem_list_elem {
-} ucp_device_mem_list_elem_t;
-
+#define UCP_DEVICE_MEM_LIST_MAX_EPS 2
 
 /**
  * @ingroup UCP_DEVICE
@@ -26,50 +25,36 @@ typedef struct ucp_mem_list_elem {
  */
 typedef struct {
     /**
-     * Allow runtime ABI compatibility checks, between host and device code.
+     * Structure version. Allow runtime ABI compatibility checks between host
+     * and device code.
      */
-    int                        version;
+    uint16_t        version;
 
     /**
      * Protocol index computed by host handle management functions when
      * creating handle.
      */
-    int                        proto_idx;
+    uint8_t         proto_idx;
 
     /**
-     * Array of pointers to UCT exported endpoints, used for multi-lane
-     * transfers.
+     * Number of UCT device endpoints found in @a uct_ep array.
      */
-    uct_ep_h                   *uct_ep;
-
-    /**
-     * Number of UCT exported endpoints found in @a uct_ep array.
-     */
-    unsigned                   num_uct_eps;
+    uint8_t         num_uct_eps;
 
     /**
      * Number of entries in the memory descriptors array @a elems.
      */
-    unsigned                   mem_list_length;
+    uint32_t        mem_list_length;
 
     /**
-     * Array of memory descriptors containing memory pairs to be used by device
-     * functions for memory transfers.
+     * Array of pointers to UCT device endpoints, used for multi-lane
+     * transfers.
      */
-    ucp_device_mem_list_elem_t elems[];
-} ucp_device_mem_list_handle_t;
+    uct_device_ep_h uct_device_eps[UCP_DEVICE_MEM_LIST_MAX_EPS];
 
-typedef ucp_device_mem_list_handle_t *ucp_device_mem_list_handle_h;
-
-
-/**
- * @ingroup UCP_DEVICE
- * @brief GPU request descriptor of a given batch
- *
- * This request tracks a batch of memory operations in progress. It can be used
- * with @ref ucp_device_progress_req to detect request completion.
- */
-typedef struct ucp_device_request {
-} ucp_device_request_t;
+    /* (mem_list_length * num_uct_eps) uct_device_mem_element objects will
+       follow this structure. The size of each element is according to the
+       selected transport. */
+} ucp_device_mem_list_handle_t;
 
 #endif /* UCP_DEVICE_TYPES_H */

src/ucp/api/device/ucp_host.h

@@ -118,7 +118,6 @@ typedef struct ucp_mem_list_params {
  *
  * Host side does not have access to the content of this descriptor.
  */
-struct ucp_device_mem_list_handle;
 typedef struct ucp_device_mem_list_handle *ucp_device_mem_list_handle_h;
 
 

src/uct/Makefile.am

@@ -20,7 +20,9 @@ nobase_dist_libuct_la_HEADERS = \
 	api/tl.h \
 	api/uct_def.h \
 	api/uct.h \
-	api/version.h
+	api/version.h \
+	api/device/uct_device_impl.h \
+	api/device/uct_device_types.h
 
 noinst_HEADERS = \
 	base/uct_md.h \

src/uct/api/device/uct_device_impl.h

@@ -0,0 +1,70 @@
+/**
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
+ *
+ * See file LICENSE for terms.
+ */
+
+#ifndef UCT_DEVICE_IMPL_H
+#define UCT_DEVICE_IMPL_H
+
+#include "uct_device_types.h"
+
+#include <uct/api/uct_def.h>
+#include <ucs/sys/compiler_def.h>
+
+
+/**
+ * @ingroup UCT_DEVICE
+ * @brief Posts one memory put operation.
+ *
+ * This device routine writes a single memory block from the local address @a address
+ * to the remote address @a remote_address using the device endpoint @a device_ep.
+ * The memory element @a mem_elem must be valid and contain the local and remote
+ * memory regions to be transferred.
+ *
+ * User can pass @a comp to track execution and completion status.
+ * The @a flags parameter can be used to modify the behavior
+ * of the routine.
+ *
+ * @param [in]  device_ep       Device endpoint to be used for the operation.
+ * @param [in]  mem_elem        Memory element representing the memory to be transferred.
+ * @param [in]  address         Local virtual address to send data from.
+ * @param [in]  remote_address  Remote virtual address to write data to.
+ * @param [in]  length          Length in bytes of the data to send.
+ * @param [in]  flags           Flags to modify the function behavior.
+ * @param [in]  comp            Completion object to track the progress of operation.
+ *
+ * @return Error code as defined by @ref ucs_status_t
+ */
+template<uct_device_level_t level = UCT_DEVICE_LEVEL_THREAD>
+UCS_F_DEVICE ucs_status_t uct_device_ep_put_single(
+        uct_device_ep_h device_ep, const uct_device_mem_element_t *mem_elem,
+        const void *address, uint64_t remote_address, size_t length,
+        uint64_t flags, uct_device_completion_t *comp)
+{
+    if (device_ep->uct_tl_id == UCT_DEVICE_TL_RC_MLX5_GDA) {
+        // return uct_rc_mlx5_gda_ep_put_single(device_ep, mem_elem, address,
+        //                                      remote_address, length, flags,
+        //                                      comp);
+    } else if (device_ep->uct_tl_id == UCT_DEVICE_TL_CUDA_IPC) {
+        // return uct_cuda_ipc_ep_put_single(device_ep, mem_elem, address,
+        //                                   remote_address, length, flags, comp);
+    }
+    return UCS_ERR_UNSUPPORTED;
+}
+
+
+/**
+ * @ingroup UCT_DEVICE
+ * @brief Initialize a device completion object.
+ *
+ * @param [out] comp  Device completion object to initialize.
+ */
+template<uct_device_level_t level = UCT_DEVICE_LEVEL_THREAD>
+UCS_F_DEVICE void uct_device_completion_init(uct_device_completion_t *comp)
+{
+    comp->count  = 0;
+    comp->status = UCS_OK;
+}
+
+#endif

src/uct/api/device/uct_device_types.h

@@ -0,0 +1,55 @@
+/**
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
+ *
+ * See file LICENSE for terms.
+ */
+
+#ifndef UCT_DEVICE_TYPES_H
+#define UCT_DEVICE_TYPES_H
+
+#include <ucs/type/status.h>
+#include <stdint.h>
+
+/**
+ * @defgroup UCT_DEVICE Device API
+ * @ingroup UCT_API
+ * * This section describes UCT Device API.
+ * @{
+ * @}
+ */
+
+/* Cooperation level when calling device functions */
+typedef enum {
+    UCT_DEVICE_LEVEL_THREAD,
+    UCT_DEVICE_LEVEL_WARP,
+    UCT_DEVICE_LEVEL_BLOCK,
+    UCT_DEVICE_LEVEL_GRID
+} uct_device_level_t;
+
+
+/* Device transport id (for internal use) */
+typedef enum {
+    UCT_DEVICE_TL_RC_MLX5_GDA,
+    UCT_DEVICE_TL_CUDA_IPC,
+    UCT_DEVICE_TL_LAST
+} uct_device_tl_id_t;
+
+
+/* Base class for all device endpoints */
+typedef struct uct_device_ep {
+    uint8_t uct_tl_id; /* Defined in uct_device_tl_id_t */
+} uct_device_ep_t;
+
+
+/* Completion object for device operations */
+typedef struct uct_device_completion {
+    uint32_t     count;  /* How many operations are pending */
+    ucs_status_t status; /* Status of the operation */
+} uct_device_completion_t;
+
+
+/* Base structure for all device memory elements */
+struct uct_device_mem_element {
+};
+
+#endif

src/uct/api/uct.h

@@ -437,6 +437,7 @@ typedef enum uct_atomic_op {
 
         /* Interface capability */
 #define UCT_IFACE_FLAG_INTER_NODE      UCS_BIT(54) /**< Interface is inter-node capable */
+#define UCT_IFACE_FLAG_DEVICE_EP       UCS_BIT(55) /**< Interface supports device endpoint */
 /**
  * @}
  */