UCT/API: introduce sys_dev field for mem_alloc

src/uct/api/uct.h

@@ -1712,6 +1712,7 @@ typedef struct uct_allocated_memory {
     ucs_memory_type_t        mem_type; /**< type of allocated memory */
     uct_md_h                 md;       /**< if method==MD: MD used to allocate the memory */
     uct_mem_h                memh;     /**< if method==MD: MD memory handle */
+    ucs_sys_device_t         sys_dev;  /**< System device for allocated memory */
 } uct_allocated_memory_t;
 
 
@@ -2501,7 +2502,10 @@ typedef enum {
     UCT_MEM_ALLOC_PARAM_FIELD_MDS            = UCS_BIT(3),
 
     /** Enables @ref uct_mem_alloc_params_t::name */
-    UCT_MEM_ALLOC_PARAM_FIELD_NAME           = UCS_BIT(4)
+    UCT_MEM_ALLOC_PARAM_FIELD_NAME           = UCS_BIT(4),
+
+    /** Enables @ref uct_mem_alloc_params_t::sys_device */
+    UCT_MEM_ALLOC_PARAM_FIELD_SYS_DEV        = UCS_BIT(5)
 } uct_mem_alloc_params_field_t;
 
 
@@ -2564,6 +2568,12 @@ typedef struct {
      * "anonymous-uct_mem_alloc" is used by default.
      */
     const char                   *name;
+
+    /**
+     * System device on which memory is to be allocated, or
+     * UCS_SYS_DEVICE_ID_UNKNOWN to allow allocating on any device.
+     */
+    ucs_sys_device_t             sys_dev;
 } uct_mem_alloc_params_t;
 
 

src/uct/base/uct_md.c

@@ -539,10 +539,11 @@ ucs_status_t uct_mem_alloc_check_params(size_t length,
 
 ucs_status_t uct_md_mem_alloc(uct_md_h md, size_t *length_p, void **address_p,
                               ucs_memory_type_t mem_type, unsigned flags,
-                              const char *alloc_name, uct_mem_h *memh_p)
+                              const char *alloc_name, ucs_sys_device_t sys_dev,
+                              uct_mem_h *memh_p)
 {
     return md->ops->mem_alloc(md, length_p, address_p, mem_type, flags,
-                              alloc_name, memh_p);
+                              alloc_name, sys_dev, memh_p);
 }
 
 ucs_status_t uct_md_mem_free(uct_md_h md, uct_mem_h memh)

src/uct/base/uct_md.h

@@ -82,6 +82,7 @@ typedef ucs_status_t (*uct_md_mem_alloc_func_t)(uct_md_h md,
                                                 ucs_memory_type_t mem_type,
                                                 unsigned flags,
                                                 const char *alloc_name,
+                                                ucs_sys_device_t sys_dev,
                                                 uct_mem_h *memh_p);
 
 typedef ucs_status_t (*uct_md_mem_free_func_t)(uct_md_h md, uct_mem_h memh);
@@ -192,11 +193,13 @@ uct_md_query_empty_md_resource(uct_md_resource_desc_t **resources_p,
  * @param [in]     flags       Memory allocation flags, see @ref uct_md_mem_flags.
  * @param [in]     name        Name of the allocated region, used to track memory
  *                             usage for debugging and profiling.
+ * @param [in]     sys_dev     System device for the allocation.
  * @param [out]    memh_p      Filled with handle for allocated region.
  */
 ucs_status_t uct_md_mem_alloc(uct_md_h md, size_t *length_p, void **address_p,
                               ucs_memory_type_t mem_type, unsigned flags,
-                              const char *alloc_name, uct_mem_h *memh_p);
+                              const char *alloc_name, ucs_sys_device_t sys_dev,
+                              uct_mem_h *memh_p);
 
 /**
  * @ingroup UCT_MD

src/uct/base/uct_mem.c

@@ -71,6 +71,7 @@ ucs_status_t uct_mem_alloc(size_t length, const uct_alloc_method_t *methods,
     const uct_alloc_method_t *method;
     ucs_log_level_t log_level;
     ucs_memory_type_t mem_type;
+    ucs_sys_device_t sys_dev;
     uct_md_attr_t md_attr;
     ucs_status_t status;
     unsigned flags;
@@ -101,12 +102,15 @@ ucs_status_t uct_mem_alloc(size_t length, const uct_alloc_method_t *methods,
                    params->name : "anonymous-uct_mem_alloc";
     mem_type     = (params->field_mask & UCT_MEM_ALLOC_PARAM_FIELD_MEM_TYPE) ?
                    params->mem_type : UCS_MEMORY_TYPE_HOST;
+    sys_dev      = (params->field_mask & UCT_MEM_ALLOC_PARAM_FIELD_SYS_DEV) ?
+                   params->sys_dev : UCS_SYS_DEVICE_ID_UNKNOWN;
     alloc_length = length;
     log_level    = (flags & UCT_MD_MEM_FLAG_HIDE_ERRORS) ? UCS_LOG_LEVEL_DEBUG :
                    UCS_LOG_LEVEL_ERROR;
 
-    ucs_trace("allocating %s: %s memory length %zu flags 0x%x", alloc_name,
-              ucs_memory_type_names[mem_type], alloc_length, flags);
+    ucs_trace("allocating %s: %s memory length %zu sys_dev %s flags 0x%x",
+              alloc_name, ucs_memory_type_names[mem_type], alloc_length,
+              ucs_topo_sys_device_get_name(sys_dev), flags);
     ucs_log_indent(1);
 
     for (method = methods; method < methods + num_methods; ++method) {
@@ -151,9 +155,8 @@ ucs_status_t uct_mem_alloc(size_t length, const uct_alloc_method_t *methods,
                  * fall-back, because this MD already exposed support for memory
                  * allocation.
                  */
-                status = uct_md_mem_alloc(md, &alloc_length, &address,
-                                          mem_type, flags, alloc_name,
-                                          &memh);
+                status = uct_md_mem_alloc(md, &alloc_length, &address, mem_type,
+                                          flags, alloc_name, sys_dev, &memh);
                 if (status != UCS_OK) {
                     ucs_log(log_level,
                             "failed to allocate %zu bytes using md %s for %s: %s",
@@ -165,6 +168,7 @@ ucs_status_t uct_mem_alloc(size_t length, const uct_alloc_method_t *methods,
                 ucs_assert(memh != UCT_MEM_HANDLE_NULL);
                 mem->md       = md;
                 mem->mem_type = mem_type;
+                mem->sys_dev  = sys_dev;
                 mem->memh     = memh;
                 goto allocated;
             }

src/uct/cuda/cuda_copy/cuda_copy_md.c

@@ -360,7 +360,8 @@ static void uct_cuda_copy_sync_memops(uct_cuda_copy_md_t *md,
 static ucs_status_t
 uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p,
                         ucs_memory_type_t mem_type, unsigned flags,
-                        const char *alloc_name, uct_mem_h *memh_p)
+                        const char *alloc_name, ucs_sys_device_t sys_dev,
+                        uct_mem_h *memh_p)
 {
     uct_cuda_copy_md_t *md = ucs_derived_of(uct_md, uct_cuda_copy_md_t);
     ucs_status_t status;

src/uct/ib/mlx5/dv/ib_mlx5dv_md.c

@@ -1986,7 +1986,7 @@ ucs_status_t
 uct_ib_mlx5_devx_device_mem_alloc(uct_md_h uct_md, size_t *length_p,
                                   void **address_p, ucs_memory_type_t mem_type,
                                   unsigned flags, const char *alloc_name,
-                                  uct_mem_h *memh_p)
+                                  ucs_sys_device_t sys_dev, uct_mem_h *memh_p)
 {
 #if HAVE_IBV_DM
     uct_ib_md_t *md           = ucs_derived_of(uct_md, uct_ib_md_t);

src/uct/rocm/copy/rocm_copy_md.c

@@ -240,7 +240,8 @@ static void uct_rocm_copy_md_close(uct_md_h uct_md) {
 static ucs_status_t
 uct_rocm_copy_mem_alloc(uct_md_h md, size_t *length_p, void **address_p,
                         ucs_memory_type_t mem_type, unsigned flags,
-                        const char *alloc_name, uct_mem_h *memh_p)
+                        const char *alloc_name, ucs_sys_device_t sys_dev,
+                        uct_mem_h *memh_p)
 {
     ucs_status_t status;
     hsa_status_t hsa_status;

src/uct/sm/mm/posix/mm_posix.c

@@ -516,7 +516,8 @@ uct_posix_segment_open(uct_mm_md_t *md, uct_mm_seg_id_t *seg_id_p, int *fd_p)
 static ucs_status_t
 uct_posix_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p,
                     ucs_memory_type_t mem_type, unsigned flags,
-                    const char *alloc_name, uct_mem_h *memh_p)
+                    const char *alloc_name, ucs_sys_device_t sys_dev,
+                    uct_mem_h *memh_p)
 {
     uct_mm_md_t                     *md = ucs_derived_of(tl_md, uct_mm_md_t);
     uct_posix_md_config_t *posix_config = ucs_derived_of(md->config,

src/uct/sm/mm/sysv/mm_sysv.c

@@ -69,7 +69,8 @@ static ucs_status_t uct_sysv_mem_attach_common(int shmid, void **address_p)
 static ucs_status_t
 uct_sysv_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p,
                    ucs_memory_type_t mem_type, unsigned flags,
-                   const char *alloc_name, uct_mem_h *memh_p)
+                   const char *alloc_name, ucs_sys_device_t sys_dev,
+                   uct_mem_h *memh_p)
 {
     uct_mm_md_t *md = ucs_derived_of(tl_md, uct_mm_md_t);
     ucs_status_t status;

src/uct/ze/copy/ze_copy_md.c

@@ -63,7 +63,8 @@ static ucs_status_t uct_ze_copy_md_query(uct_md_h md, uct_md_attr_v2_t *md_attr)
 static ucs_status_t
 uct_ze_copy_mem_alloc(uct_md_h tl_md, size_t *length_p, void **address_p,
                       ucs_memory_type_t mem_type, unsigned flags,
-                      const char *alloc_name, uct_mem_h *memh_p)
+                      const char *alloc_name, ucs_sys_device_t sys_dev,
+                      uct_mem_h *memh_p)
 {
     uct_ze_copy_md_t *md = ucs_derived_of(tl_md, uct_ze_copy_md_t);
     ze_host_mem_alloc_desc_t host_desc  = {};