From b9729a0f4f6c22159310226d40bf50854a45327c Mon Sep 17 00:00:00 2001 From: Akshay Venkatesh Date: Thu, 23 Jan 2025 21:50:45 +0000 Subject: [PATCH 1/4] UCT/API: introduce sys_dev field for mem_alloc --- src/uct/api/uct.h | 13 ++++++++++++- src/uct/base/uct_md.c | 5 +++-- src/uct/base/uct_md.h | 5 ++++- src/uct/base/uct_mem.c | 11 ++++++++--- src/uct/cuda/cuda_copy/cuda_copy_md.c | 3 ++- src/uct/ib/mlx5/dv/ib_mlx5dv_md.c | 2 +- src/uct/rocm/copy/rocm_copy_md.c | 3 ++- src/uct/sm/mm/posix/mm_posix.c | 3 ++- src/uct/sm/mm/sysv/mm_sysv.c | 3 ++- src/uct/ze/copy/ze_copy_md.c | 3 ++- 10 files changed, 38 insertions(+), 13 deletions(-) diff --git a/src/uct/api/uct.h b/src/uct/api/uct.h index 123fc2a2cfc..d8c69c6f3cc 100644 --- a/src/uct/api/uct.h +++ b/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_DEVICE = UCS_BIT(5) } uct_mem_alloc_params_field_t; @@ -2564,6 +2568,13 @@ typedef struct { * "anonymous-uct_mem_alloc" is used by default. */ const char *name; + + /** + * Index of the system device on which memory is to be allocated. + * Eg: UCS_SYS_DEVICE_ID_UNKNOWN to allocate on host memory, or a specfic + * index to allocate memory on GPU device. + */ + ucs_sys_device_t sys_dev; } uct_mem_alloc_params_t; diff --git a/src/uct/base/uct_md.c b/src/uct/base/uct_md.c index 61c0dfff10e..858bed6d98a 100644 --- a/src/uct/base/uct_md.c +++ b/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) diff --git a/src/uct/base/uct_md.h b/src/uct/base/uct_md.h index 4a75ed51bcb..bb0e16d0130 100644 --- a/src/uct/base/uct_md.h +++ b/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 diff --git a/src/uct/base/uct_mem.c b/src/uct/base/uct_mem.c index 687883bf32d..dd718fe41dd 100644 --- a/src/uct/base/uct_mem.c +++ b/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_DEVICE) ? + 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) { @@ -153,7 +157,7 @@ ucs_status_t uct_mem_alloc(size_t length, const uct_alloc_method_t *methods, */ status = uct_md_mem_alloc(md, &alloc_length, &address, mem_type, flags, alloc_name, - &memh); + sys_dev, &memh); if (status != UCS_OK) { ucs_log(log_level, "failed to allocate %zu bytes using md %s for %s: %s", @@ -165,6 +169,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; } diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index ba873106a9e..e14750eb48d 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/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; diff --git a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c b/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c index e5626040f79..b777e1f15fa 100644 --- a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c +++ b/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); diff --git a/src/uct/rocm/copy/rocm_copy_md.c b/src/uct/rocm/copy/rocm_copy_md.c index b400c7097e5..16f50b75184 100644 --- a/src/uct/rocm/copy/rocm_copy_md.c +++ b/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; diff --git a/src/uct/sm/mm/posix/mm_posix.c b/src/uct/sm/mm/posix/mm_posix.c index 667a0e85830..847979d9c52 100644 --- a/src/uct/sm/mm/posix/mm_posix.c +++ b/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, diff --git a/src/uct/sm/mm/sysv/mm_sysv.c b/src/uct/sm/mm/sysv/mm_sysv.c index dd3ba458792..9e49dccc9fc 100644 --- a/src/uct/sm/mm/sysv/mm_sysv.c +++ b/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; diff --git a/src/uct/ze/copy/ze_copy_md.c b/src/uct/ze/copy/ze_copy_md.c index f4599d727e1..0623183bee9 100644 --- a/src/uct/ze/copy/ze_copy_md.c +++ b/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 = {}; From 8c9e1a466d7c20cc316c726e8b1b06fd943f9763 Mon Sep 17 00:00:00 2001 From: Akshay Venkatesh Date: Thu, 23 Jan 2025 22:09:58 +0000 Subject: [PATCH 2/4] UCT/API: typo fix --- src/uct/api/uct.h | 2 +- src/uct/base/uct_md.h | 2 +- src/uct/base/uct_mem.c | 5 ++--- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/uct/api/uct.h b/src/uct/api/uct.h index d8c69c6f3cc..ff19a66bf07 100644 --- a/src/uct/api/uct.h +++ b/src/uct/api/uct.h @@ -2571,7 +2571,7 @@ typedef struct { /** * Index of the system device on which memory is to be allocated. - * Eg: UCS_SYS_DEVICE_ID_UNKNOWN to allocate on host memory, or a specfic + * Eg: UCS_SYS_DEVICE_ID_UNKNOWN to allocate on host memory, or a specific * index to allocate memory on GPU device. */ ucs_sys_device_t sys_dev; diff --git a/src/uct/base/uct_md.h b/src/uct/base/uct_md.h index bb0e16d0130..ed03808d24e 100644 --- a/src/uct/base/uct_md.h +++ b/src/uct/base/uct_md.h @@ -198,7 +198,7 @@ uct_md_query_empty_md_resource(uct_md_resource_desc_t **resources_p, */ 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, ucs_sys_device_t sys_dev, + const char *alloc_name, ucs_sys_device_t sys_dev, uct_mem_h *memh_p); /** diff --git a/src/uct/base/uct_mem.c b/src/uct/base/uct_mem.c index dd718fe41dd..9835c380203 100644 --- a/src/uct/base/uct_mem.c +++ b/src/uct/base/uct_mem.c @@ -155,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, - sys_dev, &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", From 61ffa4eadaf0ef82b528f13a4a9f6c01180e3110 Mon Sep 17 00:00:00 2001 From: Akshay Venkatesh Date: Mon, 27 Jan 2025 21:50:33 +0000 Subject: [PATCH 3/4] UCT/API: unify field name --- src/uct/api/uct.h | 7 +++---- src/uct/base/uct_mem.c | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/uct/api/uct.h b/src/uct/api/uct.h index ff19a66bf07..8892f8496e8 100644 --- a/src/uct/api/uct.h +++ b/src/uct/api/uct.h @@ -2505,7 +2505,7 @@ typedef enum { UCT_MEM_ALLOC_PARAM_FIELD_NAME = UCS_BIT(4), /** Enables @ref uct_mem_alloc_params_t::sys_device */ - UCT_MEM_ALLOC_PARAM_FIELD_SYS_DEVICE = UCS_BIT(5) + UCT_MEM_ALLOC_PARAM_FIELD_SYS_DEV = UCS_BIT(5) } uct_mem_alloc_params_field_t; @@ -2570,9 +2570,8 @@ typedef struct { const char *name; /** - * Index of the system device on which memory is to be allocated. - * Eg: UCS_SYS_DEVICE_ID_UNKNOWN to allocate on host memory, or a specific - * index to allocate memory on GPU device. + * 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; diff --git a/src/uct/base/uct_mem.c b/src/uct/base/uct_mem.c index 9835c380203..c08d693f79b 100644 --- a/src/uct/base/uct_mem.c +++ b/src/uct/base/uct_mem.c @@ -102,7 +102,7 @@ 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_DEVICE) ? + 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 : From 8417ce7e22f68060739a08b42aed1986105b64e5 Mon Sep 17 00:00:00 2001 From: Akshay Venkatesh Date: Fri, 31 Jan 2025 00:12:19 +0000 Subject: [PATCH 4/4] UCT/API: fix ib fxn signature; add gtest boilerplate --- src/uct/ib/mlx5/dv/ib_mlx5dv_md.c | 2 +- src/uct/ib/mlx5/ib_mlx5.h | 2 +- test/gtest/uct/test_mem.cc | 4 +++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c b/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c index b777e1f15fa..1ee41ce2bac 100644 --- a/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c +++ b/src/uct/ib/mlx5/dv/ib_mlx5dv_md.c @@ -2144,7 +2144,7 @@ static void uct_ib_mlx5dv_check_dm_ksm_reg(uct_ib_mlx5_md_t *md) status = uct_ib_mlx5_devx_device_mem_alloc(uct_md, &length, &address, UCS_MEMORY_TYPE_RDMA, 0, "check dm ksm registration", - &memh); + UCS_SYS_DEVICE_ID_UNKNOWN, &memh); if (status != UCS_OK) { ucs_debug("%s: KSM over device memory is not supported", ucs_status_string(status)); diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h index 2e082ef5e50..2dfc57de79d 100644 --- a/src/uct/ib/mlx5/ib_mlx5.h +++ b/src/uct/ib/mlx5/ib_mlx5.h @@ -1160,7 +1160,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); ucs_status_t uct_ib_mlx5_devx_device_mem_free(uct_md_h uct_md, uct_mem_h tl_memh); diff --git a/test/gtest/uct/test_mem.cc b/test/gtest/uct/test_mem.cc index 1c335ae4814..d441d18070d 100644 --- a/test/gtest/uct/test_mem.cc +++ b/test/gtest/uct/test_mem.cc @@ -87,9 +87,11 @@ UCS_TEST_P(test_mem, md_alloc) { UCT_MEM_ALLOC_PARAM_FIELD_ADDRESS | UCT_MEM_ALLOC_PARAM_FIELD_MEM_TYPE | UCT_MEM_ALLOC_PARAM_FIELD_MDS | - UCT_MEM_ALLOC_PARAM_FIELD_NAME; + UCT_MEM_ALLOC_PARAM_FIELD_NAME | + UCT_MEM_ALLOC_PARAM_FIELD_SYS_DEV; params.name = "test"; params.mem_type = UCS_MEMORY_TYPE_HOST; + params.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; /* ignored by MDs now */ params.address = address; params.mds.count = 1;