diff --git a/contrib/test_jenkins.sh b/contrib/test_jenkins.sh index 5569bcbb675..3ccfa77e9aa 100755 --- a/contrib/test_jenkins.sh +++ b/contrib/test_jenkins.sh @@ -690,6 +690,17 @@ run_ucx_perftest() { echo "==== Running ucx_perf with cuda memory ====" + if $ucx_perftest -h 2>&1 | grep -q "cuda-async" + then + echo "==== Running ucx_perf with cuda-async memory ====" + cuda_async_test_args="-t tag_lat -D contig,contig" + cuda_async_test_args+=" -m cuda-async,cuda-async -s 8 -n 10 -w 1" + run_client_server_app "$ucx_perftest" "$cuda_async_test_args" \ + "$(hostname)" 0 0 + else + echo "==== cuda-async memory is not supported, skipping ====" + fi + for memtype_cache in y n do for gdr in $gdr_options diff --git a/src/tools/perf/api/libperf.h b/src/tools/perf/api/libperf.h index 806ad9961c6..a9907ca9323 100644 --- a/src/tools/perf/api/libperf.h +++ b/src/tools/perf/api/libperf.h @@ -84,6 +84,9 @@ typedef enum { } ucx_perf_channel_mode_t; +#define UCX_PERF_ALLOC_NAME_MAX 32 /* Max memory allocator name length */ + + enum ucx_perf_test_flags { UCX_PERF_TEST_FLAG_VALIDATE = UCS_BIT(1), /* Validate data. Affects performance. */ UCX_PERF_TEST_FLAG_ONE_SIDED = UCS_BIT(2), /* For tests which involves only one side, @@ -272,6 +275,10 @@ typedef struct ucx_perf_params { ucx_perf_wait_mode_t wait_mode; /* How to wait */ ucs_memory_type_t send_mem_type; /* Send memory type */ ucs_memory_type_t recv_mem_type; /* Recv memory type */ + /* Send memory allocator name */ + char send_mem_alloc_name[UCX_PERF_ALLOC_NAME_MAX]; + /* Recv memory allocator name */ + char recv_mem_alloc_name[UCX_PERF_ALLOC_NAME_MAX]; ucx_perf_accel_dev_t send_device; /* Send memory device */ ucx_perf_accel_dev_t recv_device; /* Recv memory device */ ucs_device_level_t device_level; /* Device level */ @@ -338,10 +345,6 @@ typedef struct { } ucp_perf_daemon_req_t; -/* Allocators for each memory type */ -typedef struct ucx_perf_allocator ucx_perf_allocator_t; -extern const ucx_perf_allocator_t* ucx_perf_mem_type_allocators[]; - /* Device test dispatchers for each memory type */ typedef struct ucx_perf_device_dispatcher ucx_perf_device_dispatcher_t; extern const ucx_perf_device_dispatcher_t* ucx_perf_mem_type_device_dispatchers[]; diff --git a/src/tools/perf/cuda/cuda_alloc.c b/src/tools/perf/cuda/cuda_alloc.c index d05e56ba969..1e8bd5f1d6f 100644 --- a/src/tools/perf/cuda/cuda_alloc.c +++ b/src/tools/perf/cuda/cuda_alloc.c @@ -15,6 +15,8 @@ #include #include +#include + static ucs_status_t ucx_perf_cuda_init(ucx_perf_context_t *perf) { @@ -49,9 +51,9 @@ static ucs_status_t ucx_perf_cuda_init(ucx_perf_context_t *perf) return UCS_OK; } -static inline ucs_status_t ucx_perf_cuda_alloc(size_t length, - ucs_memory_type_t mem_type, - void **address_p) +static inline ucs_status_t ucx_perf_cuda_mem_alloc(size_t length, + ucs_memory_type_t mem_type, + void **address_p) { if (mem_type == UCS_MEMORY_TYPE_CUDA) { CUDA_CALL_RET(UCS_ERR_NO_MEMORY, cudaMalloc, address_p, length); @@ -67,12 +69,58 @@ static inline ucs_status_t ucx_perf_cuda_alloc(size_t length, return UCS_OK; } -static inline ucs_status_t -uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf, - size_t length, - ucs_memory_type_t mem_type, - unsigned flags, - uct_allocated_memory_t *alloc_mem) +static ucs_status_t ucx_perf_cuda_alloc_mem_alloc( + const ucx_perf_context_t *perf, size_t length, void **address_p) +{ + (void)perf; + return ucx_perf_cuda_mem_alloc(length, UCS_MEMORY_TYPE_CUDA, address_p); +} + +static void ucx_perf_cuda_alloc_mem_free(const ucx_perf_context_t *perf, + void *address) +{ + (void)perf; + CUDA_CALL_WARN(cudaFree, address); +} + +#if CUDART_VERSION >= 11020 +static ucs_status_t ucx_perf_cuda_async_mem_alloc( + const ucx_perf_context_t *perf, size_t length, void **address_p) +{ + cudaError_t cerr; + void *address; + + (void)perf; + + CUDA_CALL_RET(UCS_ERR_NO_MEMORY, cudaMallocAsync, &address, length, 0); + cerr = cudaStreamSynchronize(0); + if (cerr != cudaSuccess) { + ucs_error("cudaStreamSynchronize() failed: %d (%s)", (int)cerr, + cudaGetErrorString(cerr)); + CUDA_CALL_WARN(cudaFreeAsync, address, 0); + CUDA_CALL_WARN(cudaStreamSynchronize, 0); + return UCS_ERR_IO_ERROR; + } + + *address_p = address; + return UCS_OK; +} + +static void ucx_perf_cuda_async_mem_free( + const ucx_perf_context_t *perf, void *address) +{ + (void)perf; + + CUDA_CALL_WARN(cudaStreamSynchronize, 0); + CUDA_CALL_WARN(cudaFreeAsync, address, 0); + CUDA_CALL_WARN(cudaStreamSynchronize, 0); +} +#endif + +static ucs_status_t ucx_perf_cuda_uct_reg_mem( + const ucx_perf_context_t *perf, size_t length, + ucs_memory_type_t mem_type, unsigned flags, + uct_allocated_memory_t *alloc_mem) { uct_md_attr_v2_t md_attr = {.field_mask = UCT_MD_ATTR_FIELD_REG_ALIGNMENT}; size_t reg_length; @@ -85,11 +133,6 @@ uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf, return status; } - status = ucx_perf_cuda_alloc(length, mem_type, &alloc_mem->address); - if (status != UCS_OK) { - return status; - } - /* Register memory respecting MD reg_alignment */ reg_address = alloc_mem->address; reg_length = length; @@ -98,7 +141,6 @@ uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf, status = uct_perf_md_mem_reg(perf->uct.md, reg_address, reg_length, flags, mem_type, &alloc_mem->memh); if (status != UCS_OK) { - cudaFree(alloc_mem->address); ucs_error("failed to register memory"); return status; } @@ -109,24 +151,177 @@ uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf, return UCS_OK; } -static ucs_status_t uct_perf_cuda_alloc(const ucx_perf_context_t *perf, - size_t length, unsigned flags, - uct_allocated_memory_t *alloc_mem) +static ucs_status_t ucx_perf_cuda_uct_alloc_reg_mem( + const ucx_perf_context_t *perf, size_t length, + ucs_memory_type_t mem_type, unsigned flags, + uct_allocated_memory_t *alloc_mem) { - return uct_perf_cuda_alloc_reg_mem(perf, length, UCS_MEMORY_TYPE_CUDA, - flags, alloc_mem); + ucs_status_t status; + + status = ucx_perf_cuda_mem_alloc(length, mem_type, &alloc_mem->address); + if (status != UCS_OK) { + return status; + } + + status = ucx_perf_cuda_uct_reg_mem(perf, length, mem_type, flags, + alloc_mem); + if (status != UCS_OK) { + CUDA_CALL_WARN(cudaFree, alloc_mem->address); + return status; + } + + return UCS_OK; +} + +static ucs_status_t ucx_perf_cuda_uct_alloc(const ucx_perf_context_t *perf, + size_t length, unsigned flags, + uct_allocated_memory_t *alloc_mem) +{ + return ucx_perf_cuda_uct_alloc_reg_mem(perf, length, UCS_MEMORY_TYPE_CUDA, + flags, alloc_mem); } -static ucs_status_t uct_perf_cuda_managed_alloc(const ucx_perf_context_t *perf, - size_t length, unsigned flags, - uct_allocated_memory_t *alloc_mem) +static ucs_status_t ucx_perf_cuda_managed_uct_alloc( + const ucx_perf_context_t *perf, size_t length, unsigned flags, + uct_allocated_memory_t *alloc_mem) { - return uct_perf_cuda_alloc_reg_mem(perf, length, UCS_MEMORY_TYPE_CUDA_MANAGED, - flags, alloc_mem); + return ucx_perf_cuda_uct_alloc_reg_mem(perf, length, + UCS_MEMORY_TYPE_CUDA_MANAGED, flags, + alloc_mem); } -static void uct_perf_cuda_free(const ucx_perf_context_t *perf, - uct_allocated_memory_t *alloc_mem) +#if CUDART_VERSION >= 11020 +/* Resolve async allocation memory type from UCX_CUDA_COPY_ASYNC_MEM_TYPE. */ +static ucs_memory_type_t ucx_perf_cuda_async_configured_mem_type(void) +{ + static int initialized = 0; + static ucs_memory_type_t cached_mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED; + uct_component_h *components = NULL; + unsigned num_components = 0; + ucs_memory_type_t result = UCS_MEMORY_TYPE_CUDA_MANAGED; + uct_component_attr_t component_attr; + uct_md_config_t *md_config; + char value[64]; + ucs_memory_type_t mem_type; + ucs_status_t status; + unsigned i; + + if (initialized) { + return cached_mem_type; + } + + status = uct_query_components(&components, &num_components); + if (status != UCS_OK) { + ucs_debug("failed to query UCT components: %s", + ucs_status_string(status)); + goto out; + } + + for (i = 0; i < num_components; ++i) { + component_attr.field_mask = UCT_COMPONENT_ATTR_FIELD_NAME; + status = uct_component_query(components[i], &component_attr); + if ((status != UCS_OK) || strcmp(component_attr.name, "cuda_cpy")) { + continue; + } + + status = uct_md_config_read(components[i], NULL, NULL, &md_config); + if (status != UCS_OK) { + ucs_debug("failed to read cuda_cpy MD config: %s", + ucs_status_string(status)); + break; + } + + status = uct_config_get(md_config, "ASYNC_MEM_TYPE", value, + sizeof(value)); + uct_config_release(md_config); + if (status != UCS_OK) { + ucs_debug("failed to get ASYNC_MEM_TYPE: %s", + ucs_status_string(status)); + break; + } + + for (mem_type = 0; mem_type < UCS_MEMORY_TYPE_LAST; ++mem_type) { + if (strcmp(value, ucs_memory_type_names[mem_type])) { + continue; + } + + if ((mem_type == UCS_MEMORY_TYPE_CUDA) || + (mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED)) { + result = mem_type; + } else { + ucs_warn("wrong memory type for async memory allocations: " + "\"%s\"; cuda-managed will be used instead", + value); + } + break; + } + break; + } + + uct_release_component_list(components); + +out: + cached_mem_type = result; + initialized = 1; + return result; +} + +static ucs_memory_type_t +ucx_perf_cuda_async_mem_type(const ucx_perf_context_t *perf, + const void *address, size_t length) +{ + ucs_memory_type_t mem_type; + ucs_status_t status; + + if (perf->params.api != UCX_PERF_API_UCT) { + return UCS_MEMORY_TYPE_LAST; + } + + status = uct_md_detect_memory_type(perf->uct.md, address, length, + &mem_type); + if (status == UCS_OK) { + return mem_type; + } + + mem_type = ucx_perf_cuda_async_configured_mem_type(); + ucs_debug("failed to detect cuda async memory type: %s, using %s", + ucs_status_string(status), ucs_memory_type_names[mem_type]); + return mem_type; +} + +static ucs_memory_type_t +ucx_perf_cuda_async_resolve_mem_type(const ucx_perf_allocator_t *allocator) +{ + (void)allocator; + return ucx_perf_cuda_async_configured_mem_type(); +} + +static ucs_status_t ucx_perf_cuda_async_uct_alloc( + const ucx_perf_context_t *perf, size_t length, unsigned flags, + uct_allocated_memory_t *alloc_mem) +{ + ucs_memory_type_t mem_type; + ucs_status_t status; + + status = ucx_perf_cuda_async_mem_alloc(perf, length, &alloc_mem->address); + if (status != UCS_OK) { + return status; + } + + mem_type = ucx_perf_cuda_async_mem_type(perf, alloc_mem->address, length); + status = ucx_perf_cuda_uct_reg_mem(perf, length, mem_type, flags, + alloc_mem); + if (status != UCS_OK) { + ucx_perf_cuda_async_mem_free(perf, alloc_mem->address); + return status; + } + + return UCS_OK; +} +#endif + +static void ucx_perf_cuda_uct_dereg(const ucx_perf_context_t *perf, + uct_allocated_memory_t *alloc_mem) { ucs_status_t status; @@ -136,10 +331,24 @@ static void uct_perf_cuda_free(const ucx_perf_context_t *perf, if (status != UCS_OK) { ucs_error("failed to deregister memory"); } +} +static void ucx_perf_cuda_uct_free(const ucx_perf_context_t *perf, + uct_allocated_memory_t *alloc_mem) +{ + ucx_perf_cuda_uct_dereg(perf, alloc_mem); CUDA_CALL_WARN(cudaFree, alloc_mem->address); } +#if CUDART_VERSION >= 11020 +static void ucx_perf_cuda_async_uct_free(const ucx_perf_context_t *perf, + uct_allocated_memory_t *alloc_mem) +{ + ucx_perf_cuda_uct_dereg(perf, alloc_mem); + ucx_perf_cuda_async_mem_free(perf, alloc_mem->address); +} +#endif + static void ucx_perf_cuda_memcpy(void *dst, ucs_memory_type_t dst_mem_type, const void *src, ucs_memory_type_t src_mem_type, size_t count) @@ -155,29 +364,71 @@ static void* ucx_perf_cuda_memset(void *dst, int value, size_t count) return dst; } +#if CUDART_VERSION >= 11020 +static ucx_perf_allocator_t cuda_async_allocator = { + .name = "cuda-async", + .default_mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED, + .init = ucx_perf_cuda_init, + .uct_alloc = ucx_perf_cuda_async_uct_alloc, + .uct_free = ucx_perf_cuda_async_uct_free, + .mem_alloc = ucx_perf_cuda_async_mem_alloc, + .mem_free = ucx_perf_cuda_async_mem_free, + .detect_mem_type = ucx_perf_cuda_async_mem_type, + .resolve_mem_type = ucx_perf_cuda_async_resolve_mem_type, + .memcpy = ucx_perf_cuda_memcpy, + .memset = ucx_perf_cuda_memset +}; +#endif + +static ucx_perf_allocator_t cuda_ucp_allocator = { + .name = "cuda", + .default_mem_type = UCS_MEMORY_TYPE_CUDA, + .init = ucx_perf_cuda_init, + .uct_alloc = ucx_perf_cuda_uct_alloc, + .uct_free = ucx_perf_cuda_uct_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_cuda_memcpy, + .memset = ucx_perf_cuda_memset +}; + +static ucx_perf_allocator_t cuda_alloc_allocator = { + .name = "cuda-alloc", + .default_mem_type = UCS_MEMORY_TYPE_CUDA, + .init = ucx_perf_cuda_init, + .uct_alloc = ucx_perf_cuda_uct_alloc, + .uct_free = ucx_perf_cuda_uct_free, + .mem_alloc = ucx_perf_cuda_alloc_mem_alloc, + .mem_free = ucx_perf_cuda_alloc_mem_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_cuda_memcpy, + .memset = ucx_perf_cuda_memset +}; + +static ucx_perf_allocator_t cuda_managed_allocator = { + .name = "cuda-managed", + .default_mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED, + .init = ucx_perf_cuda_init, + .uct_alloc = ucx_perf_cuda_managed_uct_alloc, + .uct_free = ucx_perf_cuda_uct_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_cuda_memcpy, + .memset = ucx_perf_cuda_memset +}; + UCS_STATIC_INIT { - static ucx_perf_allocator_t cuda_allocator = { - .mem_type = UCS_MEMORY_TYPE_CUDA, - .init = ucx_perf_cuda_init, - .uct_alloc = uct_perf_cuda_alloc, - .uct_free = uct_perf_cuda_free, - .memcpy = ucx_perf_cuda_memcpy, - .memset = ucx_perf_cuda_memset - }; - static ucx_perf_allocator_t cuda_managed_allocator = { - .mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED, - .init = ucx_perf_cuda_init, - .uct_alloc = uct_perf_cuda_managed_alloc, - .uct_free = uct_perf_cuda_free, - .memcpy = ucx_perf_cuda_memcpy, - .memset = ucx_perf_cuda_memset - }; - - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA] = &cuda_allocator; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA_MANAGED] = &cuda_managed_allocator; + ucx_perf_allocator_register(&cuda_ucp_allocator); + ucx_perf_allocator_register(&cuda_alloc_allocator); + ucx_perf_allocator_register(&cuda_managed_allocator); +#if CUDART_VERSION >= 11020 + ucx_perf_allocator_register(&cuda_async_allocator); +#endif } UCS_STATIC_CLEANUP { - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA] = NULL; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA_MANAGED] = NULL; +#if CUDART_VERSION >= 11020 + ucx_perf_allocator_unregister(&cuda_async_allocator); +#endif + ucx_perf_allocator_unregister(&cuda_managed_allocator); + ucx_perf_allocator_unregister(&cuda_alloc_allocator); + ucx_perf_allocator_unregister(&cuda_ucp_allocator); } diff --git a/src/tools/perf/lib/libperf.c b/src/tools/perf/lib/libperf.c index 46b00338358..2f116ba5cae 100644 --- a/src/tools/perf/lib/libperf.c +++ b/src/tools/perf/lib/libperf.c @@ -84,11 +84,25 @@ typedef struct { } ucp_perf_flush_context_t; -const ucx_perf_allocator_t* ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_LAST]; +const ucx_perf_allocator_t *ucx_perf_allocators[UCX_PERF_ALLOCATOR_MAX]; +unsigned ucx_perf_num_allocators; const ucx_perf_device_dispatcher_t* ucx_perf_mem_type_device_dispatchers[UCS_MEMORY_TYPE_LAST]; +const ucx_perf_allocator_t *ucx_perf_allocator_by_name(const char *name) +{ + unsigned i; + + for (i = 0; i < ucx_perf_num_allocators; ++i) { + if (!strcmp(name, ucx_perf_allocators[i]->name)) { + return ucx_perf_allocators[i]; + } + } + + return NULL; +} + static const char *perf_iface_ops[] = { [ucs_ilog2(UCT_IFACE_FLAG_AM_SHORT)] = "am short", [ucs_ilog2(UCT_IFACE_FLAG_AM_BCOPY)] = "am bcopy", @@ -2164,36 +2178,51 @@ ucx_perf_funcs_t ucx_perf_funcs[] = { ucp_perf_test_dispatch, ucp_perf_thread_barrier} }; +static const ucx_perf_allocator_t *ucx_perf_params_allocator( + const ucx_perf_params_t *params, int is_send) +{ + const char *alloc_name; + + alloc_name = ucx_perf_mem_alloc_name(params, is_send); + return ucx_perf_allocator_by_name(alloc_name); +} + ucs_status_t ucx_perf_allocators_init(ucx_perf_context_t *perf, const ucx_perf_params_t *params) { - ucs_debug("set send allocator by send mem type %s", - ucs_memory_type_names[params->send_mem_type]); - perf->send_allocator = ucx_perf_mem_type_allocators[params->send_mem_type]; + ucs_memory_type_t send_mem_type; + ucs_memory_type_t recv_mem_type; - ucs_debug("set recv allocator by recv mem type %s", - ucs_memory_type_names[params->recv_mem_type]); - perf->recv_allocator = ucx_perf_mem_type_allocators[params->recv_mem_type]; + perf->send_allocator = ucx_perf_params_allocator(params, 1); + perf->recv_allocator = ucx_perf_params_allocator(params, 0); if ((perf->send_allocator == NULL) || (perf->recv_allocator == NULL)) { - ucs_error("Unsupported memory types %s<->%s", - ucs_memory_type_names[params->send_mem_type], - ucs_memory_type_names[params->recv_mem_type]); + ucs_error("Unsupported memory allocators %s<->%s", + ucx_perf_mem_alloc_name(params, 1), + ucx_perf_mem_alloc_name(params, 0)); return UCS_ERR_UNSUPPORTED; } + ucs_debug("set send allocator %s", perf->send_allocator->name); + ucs_debug("set recv allocator %s", perf->recv_allocator->name); + if (perf->params.api == UCX_PERF_API_UCT) { - if (perf->send_allocator->mem_type != UCS_MEMORY_TYPE_HOST) { + send_mem_type = perf->send_allocator->resolve_mem_type( + perf->send_allocator); + recv_mem_type = perf->recv_allocator->resolve_mem_type( + perf->recv_allocator); + + if (send_mem_type != UCS_MEMORY_TYPE_HOST) { ucs_diag("UCT tests also copy one-byte value from %s memory to " "%s send memory, which may impact performance results", ucs_memory_type_names[UCS_MEMORY_TYPE_HOST], - ucs_memory_type_names[perf->send_allocator->mem_type]); + ucs_memory_type_names[send_mem_type]); } - if (perf->recv_allocator->mem_type != UCS_MEMORY_TYPE_HOST) { + if (recv_mem_type != UCS_MEMORY_TYPE_HOST) { ucs_diag("UCT tests also copy one-byte value from %s recv memory " "to %s memory, which may impact performance results", - ucs_memory_type_names[perf->recv_allocator->mem_type], + ucs_memory_type_names[recv_mem_type], ucs_memory_type_names[UCS_MEMORY_TYPE_HOST]); } } diff --git a/src/tools/perf/lib/libperf_int.h b/src/tools/perf/lib/libperf_int.h index 8ec7a6ae085..0775e8146ec 100644 --- a/src/tools/perf/lib/libperf_int.h +++ b/src/tools/perf/lib/libperf_int.h @@ -17,12 +17,14 @@ #include #endif +#include BEGIN_C_DECLS /** @file libperf_int.h */ #include +#include #include #include @@ -32,6 +34,7 @@ BEGIN_C_DECLS #define ADDR_BUF_SIZE 4096 #define EXTRA_INFO_SIZE 256 #define ONESIDED_SIGNAL_SIZE sizeof(uint64_t) +#define UCX_PERF_ALLOCATOR_MAX (UCS_MEMORY_TYPE_LAST + 16) #define UCX_PERF_TEST_FOREACH(perf) \ while (!ucx_perf_context_done(perf)) @@ -41,6 +44,7 @@ BEGIN_C_DECLS typedef struct ucx_perf_context ucx_perf_context_t; +typedef struct ucx_perf_allocator ucx_perf_allocator_t; typedef struct uct_peer uct_peer_t; typedef struct ucp_perf_request ucp_perf_request_t; typedef struct ucx_perf_thread_context ucx_perf_thread_context_t; @@ -55,6 +59,18 @@ typedef ucs_status_t (*ucx_perf_uct_alloc_func_t)( typedef void (*ucx_perf_uct_free_func_t)(const ucx_perf_context_t *perf, uct_allocated_memory_t *alloc_mem); +typedef ucs_status_t (*ucx_perf_mem_alloc_func_t)( + const ucx_perf_context_t *perf, size_t length, void **address_p); + +typedef void (*ucx_perf_mem_free_func_t)(const ucx_perf_context_t *perf, + void *address); + +typedef ucs_memory_type_t (*ucx_perf_mem_type_func_t)( + const ucx_perf_context_t *perf, const void *address, size_t length); + +typedef ucs_memory_type_t (*ucx_perf_resolve_mem_type_func_t)( + const ucx_perf_allocator_t *allocator); + typedef void (*ucx_perf_memcpy_func_t)(void *dst, ucs_memory_type_t dst_mem_type, const void *src, @@ -64,14 +80,75 @@ typedef void (*ucx_perf_memcpy_func_t)(void *dst, typedef void *(*ucx_perf_memset_func_t)(void *dst, int value, size_t count); struct ucx_perf_allocator { - ucs_memory_type_t mem_type; - ucx_perf_init_func_t init; - ucx_perf_uct_alloc_func_t uct_alloc; - ucx_perf_uct_free_func_t uct_free; - ucx_perf_memcpy_func_t memcpy; - ucx_perf_memset_func_t memset; + const char *name; + ucs_memory_type_t default_mem_type; + ucx_perf_init_func_t init; + ucx_perf_uct_alloc_func_t uct_alloc; + ucx_perf_uct_free_func_t uct_free; + ucx_perf_mem_alloc_func_t mem_alloc; + ucx_perf_mem_free_func_t mem_free; + ucx_perf_mem_type_func_t detect_mem_type; + ucx_perf_resolve_mem_type_func_t resolve_mem_type; + ucx_perf_memcpy_func_t memcpy; + ucx_perf_memset_func_t memset; }; +extern const ucx_perf_allocator_t *ucx_perf_allocators[UCX_PERF_ALLOCATOR_MAX]; +extern unsigned ucx_perf_num_allocators; + +static UCS_F_ALWAYS_INLINE ucs_status_t +ucx_perf_allocator_register(const ucx_perf_allocator_t *allocator) +{ + const char *name = allocator->name; + unsigned i; + + if (name == NULL) { + ucs_error("perftest memory allocator name is not set"); + return UCS_ERR_INVALID_PARAM; + } + + if (allocator->resolve_mem_type == NULL) { + ucs_error("perftest memory allocator resolve_mem_type is not set"); + return UCS_ERR_INVALID_PARAM; + } + + if (strlen(name) >= UCX_PERF_ALLOC_NAME_MAX) { + ucs_error("perftest memory allocator name is too long: \"%s\"", name); + return UCS_ERR_INVALID_PARAM; + } + + for (i = 0; i < ucx_perf_num_allocators; ++i) { + if (!strcmp(name, ucx_perf_allocators[i]->name)) { + return (ucx_perf_allocators[i] == allocator) ? + UCS_OK : UCS_ERR_ALREADY_EXISTS; + } + } + + if (ucx_perf_num_allocators == ucs_static_array_size(ucx_perf_allocators)) { + ucs_error("too many perftest memory allocators"); + return UCS_ERR_EXCEEDS_LIMIT; + } + + ucx_perf_allocators[ucx_perf_num_allocators++] = allocator; + return UCS_OK; +} + +static UCS_F_ALWAYS_INLINE void +ucx_perf_allocator_unregister(const ucx_perf_allocator_t *allocator) +{ + unsigned i; + + for (i = 0; i < ucx_perf_num_allocators; ++i) { + if (ucx_perf_allocators[i] == allocator) { + --ucx_perf_num_allocators; + ucx_perf_allocators[i] = + ucx_perf_allocators[ucx_perf_num_allocators]; + ucx_perf_allocators[ucx_perf_num_allocators] = NULL; + return; + } + } +} + typedef ucs_status_t (*ucp_perf_dispatch_func_t)(ucx_perf_context_t *perf); @@ -223,6 +300,43 @@ void ucx_perf_report(ucx_perf_context_t *perf); ucs_status_t ucx_perf_allocators_init_thread(ucx_perf_context_t *perf); +const ucx_perf_allocator_t *ucx_perf_allocator_by_name(const char *name); + + +static UCS_F_ALWAYS_INLINE ucs_memory_type_t +ucx_perf_allocator_default_resolve_mem_type( + const ucx_perf_allocator_t *allocator) +{ + return allocator->default_mem_type; +} + +static UCS_F_ALWAYS_INLINE ucs_memory_type_t ucx_perf_allocator_mem_type( + const ucx_perf_context_t *perf, const ucx_perf_allocator_t *allocator, + const void *address, size_t length) +{ + if (allocator->detect_mem_type != NULL) { + return allocator->detect_mem_type(perf, address, length); + } + + return allocator->resolve_mem_type(allocator); +} + +static UCS_F_ALWAYS_INLINE const char * +ucx_perf_mem_alloc_name(const ucx_perf_params_t *params, int is_send) +{ + const char *alloc_name; + ucs_memory_type_t mem_type; + + alloc_name = is_send ? params->send_mem_alloc_name : + params->recv_mem_alloc_name; + if (alloc_name[0] != '\0') { + return alloc_name; + } + + mem_type = is_send ? params->send_mem_type : params->recv_mem_type; + return ucs_memory_type_names[mem_type]; +} + static UCS_F_ALWAYS_INLINE int ucx_perf_context_done(ucx_perf_context_t *perf) { return ucs_unlikely((perf->current.iters >= perf->max_iter) || diff --git a/src/tools/perf/lib/libperf_memory.c b/src/tools/perf/lib/libperf_memory.c index 9502d6b5190..a91450621d6 100644 --- a/src/tools/perf/lib/libperf_memory.c +++ b/src/tools/perf/lib/libperf_memory.c @@ -43,42 +43,75 @@ static ucs_status_t ucp_perf_test_alloc_iov_mem(ucp_perf_datatype_t datatype, static ucs_status_t ucp_perf_mem_alloc(const ucx_perf_context_t *perf, size_t length, - ucs_memory_type_t mem_type, + const ucx_perf_allocator_t *allocator, void **address_p, ucp_mem_h *memh_p) { + int external_alloc = (allocator->mem_alloc != NULL); ucp_mem_map_params_t params; ucp_mem_attr_t attr; + void *address; ucs_status_t status; + ucs_assert((allocator->mem_alloc == NULL) == (allocator->mem_free == NULL)); params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS | UCP_MEM_MAP_PARAM_FIELD_LENGTH | - UCP_MEM_MAP_PARAM_FIELD_FLAGS | - UCP_MEM_MAP_PARAM_FIELD_MEMORY_TYPE; + UCP_MEM_MAP_PARAM_FIELD_FLAGS; params.address = NULL; - params.memory_type = mem_type; + params.memory_type = allocator->resolve_mem_type(allocator); params.length = length; - params.flags = UCP_MEM_MAP_ALLOCATE; + params.flags = external_alloc ? 0 : UCP_MEM_MAP_ALLOCATE; + if (!external_alloc || (allocator->detect_mem_type == NULL)) { + params.field_mask |= UCP_MEM_MAP_PARAM_FIELD_MEMORY_TYPE; + } + if (perf->params.flags & UCX_PERF_TEST_FLAG_MAP_NONBLOCK) { params.flags |= UCP_MEM_MAP_NONBLOCK; } + if (external_alloc) { + status = allocator->mem_alloc(perf, length, &address); + if (status != UCS_OK) { + return status; + } + + params.address = address; + params.memory_type = ucx_perf_allocator_mem_type(perf, allocator, + address, length); + if (params.memory_type != UCS_MEMORY_TYPE_LAST) { + params.field_mask |= UCP_MEM_MAP_PARAM_FIELD_MEMORY_TYPE; + } + } + status = ucp_mem_map(perf->ucp.context, ¶ms, memh_p); if (status != UCS_OK) { - return status; + goto err_free_mem; + } + + if (external_alloc) { + *address_p = address; + return UCS_OK; } attr.field_mask = UCP_MEM_ATTR_FIELD_ADDRESS; status = ucp_mem_query(*memh_p, &attr); if (status != UCS_OK) { ucp_mem_unmap(perf->ucp.context, *memh_p); - return status; + goto err_free_mem; } *address_p = attr.address; return UCS_OK; + +err_free_mem: + if (external_alloc) { + allocator->mem_free(perf, address); + } + return status; } -static void ucp_perf_mem_free(const ucx_perf_context_t *perf, ucp_mem_h memh) +static void ucp_perf_mem_free(const ucx_perf_context_t *perf, ucp_mem_h memh, + const ucx_perf_allocator_t *allocator, + void *address) { ucs_status_t status; @@ -86,6 +119,10 @@ static void ucp_perf_mem_free(const ucx_perf_context_t *perf, ucp_mem_h memh) if (status != UCS_OK) { ucs_warn("ucp_mem_unmap() failed: %s", ucs_status_string(status)); } + + if (allocator->mem_free != NULL) { + allocator->mem_free(perf, address); + } } static void ucp_perf_mem_buf_free(ucx_perf_exported_mem_t *exported_mem) @@ -114,7 +151,7 @@ ucs_status_t ucp_perf_test_alloc_mem(ucx_perf_context_t *perf) /* Allocate send buffer memory */ status = ucp_perf_mem_alloc(perf, buffer_size * params->thread_count + ONESIDED_SIGNAL_SIZE, - params->send_mem_type, &perf->send_buffer, + perf->send_allocator, &perf->send_buffer, &perf->ucp.send_memh); if (status != UCS_OK) { goto err; @@ -128,7 +165,7 @@ ucs_status_t ucp_perf_test_alloc_mem(ucx_perf_context_t *perf) /* Allocate receive buffer memory */ status = ucp_perf_mem_alloc(perf, buffer_size * params->thread_count + ONESIDED_SIGNAL_SIZE, - params->recv_mem_type, &perf->recv_buffer, + perf->recv_allocator, &perf->recv_buffer, &perf->ucp.recv_memh); if (status != UCS_OK) { goto err_free_send_buffer; @@ -189,11 +226,13 @@ ucs_status_t ucp_perf_test_alloc_mem(ucx_perf_context_t *perf) err_free_am_hdr: free(perf->ucp.am_hdr); err_free_buffers: - ucp_perf_mem_free(perf, perf->ucp.recv_memh); + ucp_perf_mem_free(perf, perf->ucp.recv_memh, perf->recv_allocator, + perf->recv_buffer); ucp_perf_mem_buf_free(&perf->ucp.recv_exported_mem); ucp_perf_mem_buf_free(&perf->ucp.send_exported_mem); err_free_send_buffer: - ucp_perf_mem_free(perf, perf->ucp.send_memh); + ucp_perf_mem_free(perf, perf->ucp.send_memh, perf->send_allocator, + perf->send_buffer); err: return UCS_ERR_NO_MEMORY; } @@ -203,8 +242,10 @@ void ucp_perf_test_free_mem(ucx_perf_context_t *perf) free(perf->ucp.recv_iov); free(perf->ucp.send_iov); free(perf->ucp.am_hdr); - ucp_perf_mem_free(perf, perf->ucp.recv_memh); - ucp_perf_mem_free(perf, perf->ucp.send_memh); + ucp_perf_mem_free(perf, perf->ucp.recv_memh, perf->recv_allocator, + perf->recv_buffer); + ucp_perf_mem_free(perf, perf->ucp.send_memh, perf->send_allocator, + perf->send_buffer); ucp_perf_mem_buf_free(&perf->ucp.recv_exported_mem); ucp_perf_mem_buf_free(&perf->ucp.send_exported_mem); } @@ -293,7 +334,6 @@ ucs_status_t uct_perf_test_alloc_mem(ucx_perf_context_t *perf) status = perf->send_allocator->uct_alloc(perf, buffer_size * params->thread_count, flags, &perf->uct.send_mem); - if (status != UCS_OK) { goto err; } @@ -344,26 +384,36 @@ void uct_perf_test_free_mem(ucx_perf_context_t *perf) void ucx_perf_global_init() { static ucx_perf_allocator_t host_allocator = { - .mem_type = UCS_MEMORY_TYPE_HOST, - .init = (ucx_perf_init_func_t)ucs_empty_function_return_success, - .uct_alloc = uct_perf_test_alloc_host, - .uct_free = uct_perf_test_free_host, - .memcpy = ucx_perf_test_memcpy_host, - .memset = memset + .name = "host", + .default_mem_type = UCS_MEMORY_TYPE_HOST, + .init = (ucx_perf_init_func_t) + ucs_empty_function_return_success, + .uct_alloc = uct_perf_test_alloc_host, + .uct_free = uct_perf_test_free_host, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_test_memcpy_host, + .memset = memset }; static ucx_perf_allocator_t rdma_allocator = { - .mem_type = UCS_MEMORY_TYPE_RDMA, - .init = (ucx_perf_init_func_t)ucs_empty_function_return_success, - .uct_alloc = (ucx_perf_uct_alloc_func_t)ucs_empty_function_do_assert, - .uct_free = (ucx_perf_uct_free_func_t)ucs_empty_function_do_assert, - .memcpy = (ucx_perf_memcpy_func_t)ucs_empty_function_do_assert, - .memset = (ucx_perf_memset_func_t)ucs_empty_function_do_assert + .name = "rdma", + .default_mem_type = UCS_MEMORY_TYPE_RDMA, + .init = (ucx_perf_init_func_t) + ucs_empty_function_return_success, + .uct_alloc = (ucx_perf_uct_alloc_func_t) + ucs_empty_function_do_assert, + .uct_free = (ucx_perf_uct_free_func_t) + ucs_empty_function_do_assert, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = (ucx_perf_memcpy_func_t) + ucs_empty_function_do_assert, + .memset = (ucx_perf_memset_func_t) + ucs_empty_function_do_assert }; UCS_MODULE_FRAMEWORK_DECLARE(ucx_perftest); - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_HOST] = &host_allocator; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_RDMA] = &rdma_allocator; + ucx_perf_allocator_register(&host_allocator); + ucx_perf_allocator_register(&rdma_allocator); /* FIXME Memtype allocator modules must be loaded to global scope, otherwise * alloc hooks, which are using dlsym() to get pointer to original function, diff --git a/src/tools/perf/lib/uct_tests.cc b/src/tools/perf/lib/uct_tests.cc index f7f44d5f213..a94cd441416 100644 --- a/src/tools/perf/lib/uct_tests.cc +++ b/src/tools/perf/lib/uct_tests.cc @@ -133,7 +133,8 @@ class uct_perf_test_runner { const void *src_sn, const ucx_perf_allocator_t *allocator) const { - if (ucs_likely(allocator->mem_type == UCS_MEMORY_TYPE_HOST)) { + if (ucs_likely(allocator->resolve_mem_type(allocator) == + UCS_MEMORY_TYPE_HOST)) { ucs_assert(dst_mem_type == UCS_MEMORY_TYPE_HOST); *reinterpret_cast(dst_sn) = *reinterpret_cast(src_sn); } diff --git a/src/tools/perf/perftest.c b/src/tools/perf/perftest.c index 5f0bfbbb072..25f820ab719 100644 --- a/src/tools/perf/perftest.c +++ b/src/tools/perf/perftest.c @@ -195,6 +195,10 @@ ucs_status_t init_test_params(perftest_params_t *params) params->super.uct.am_hdr_size = 8; params->super.send_mem_type = UCS_MEMORY_TYPE_HOST; params->super.recv_mem_type = UCS_MEMORY_TYPE_HOST; + ucs_strncpy_safe(params->super.send_mem_alloc_name, "host", + UCX_PERF_ALLOC_NAME_MAX); + ucs_strncpy_safe(params->super.recv_mem_alloc_name, "host", + UCX_PERF_ALLOC_NAME_MAX); params->super.send_device = default_dev; params->super.recv_device = default_dev; params->super.device_level = UCS_DEVICE_LEVEL_THREAD; diff --git a/src/tools/perf/perftest_params.c b/src/tools/perf/perftest_params.c index 8e2fa07f678..acd03c345e0 100644 --- a/src/tools/perf/perftest_params.c +++ b/src/tools/perf/perftest_params.c @@ -24,13 +24,16 @@ const struct option TEST_PARAMS_ARGS_LONG[] = {0, 0, 0, 0} }; -static void print_memory_type_usage(void) +static void print_memory_allocator_usage(void) { - ucs_memory_type_t it; + const ucx_perf_allocator_t *allocator; + unsigned i; - ucs_memory_type_for_each(it) { - printf(" %s - %s\n", ucs_memory_type_names[it], - ucs_memory_type_descs[it]); + for (i = 0; i < ucx_perf_num_allocators; ++i) { + allocator = ucx_perf_allocators[i]; + printf(" %s - %s\n", + allocator->name, + ucs_memory_type_descs[allocator->resolve_mem_type(allocator)]); } } @@ -90,9 +93,10 @@ static void usage(const struct perftest_context *ctx, const char *program) ctx->params.super.msg_size_list[0]); printf(" for example: \"-s 16,48,8192,8192,14\"\n"); printf(" compact form example: \"-s 1024:16 expands to [1024, ..., 1024] with 16 elements\n"); - printf(" -m [,]\n"); - printf(" memory type of message for sender and receiver (host)\n"); - print_memory_type_usage(); + printf(" -m [,]\n"); + printf(" memory allocator for sender and receiver " + "(host)\n"); + print_memory_allocator_usage(); printf(" -n number of iterations to run (%"PRIu64")\n", ctx->params.super.max_iter); printf(" -w number of warm-up iterations (%"PRIu64")\n", ctx->params.super.warmup_iter); @@ -241,6 +245,33 @@ static ucs_status_t parse_mem_type(const char *opt_arg, return UCS_ERR_INVALID_PARAM; } +static ucs_status_t parse_perf_mem_allocator(const char *opt_arg, + ucs_memory_type_t *mem_type, + char *alloc_name) +{ + const ucx_perf_allocator_t *allocator; + + if (opt_arg == NULL) { + ucs_error("memory allocator string is NULL"); + return UCS_ERR_INVALID_PARAM; + } + + if (strlen(opt_arg) >= UCX_PERF_ALLOC_NAME_MAX) { + ucs_error("memory allocator name is too long: \"%s\"", opt_arg); + return UCS_ERR_INVALID_PARAM; + } + + allocator = ucx_perf_allocator_by_name(opt_arg); + if (allocator != NULL) { + *mem_type = allocator->resolve_mem_type(allocator); + ucs_strncpy_safe(alloc_name, opt_arg, UCX_PERF_ALLOC_NAME_MAX); + return UCS_OK; + } + + ucs_error("unsupported memory allocator: \"%s\"", opt_arg); + return UCS_ERR_INVALID_PARAM; +} + static ucs_status_t parse_accel_device(char *opt_arg, ucx_perf_accel_dev_t *dev) { @@ -279,13 +310,15 @@ parse_accel_device(char *opt_arg, ucx_perf_accel_dev_t *dev) static ucs_status_t parse_mem_type_params(const char *opt_arg, ucs_memory_type_t *send_mem_type, - ucs_memory_type_t *recv_mem_type) + ucs_memory_type_t *recv_mem_type, + char *send_alloc_name, + char *recv_alloc_name) { - const char *delim = ","; - char *token = strtok((char*)opt_arg, delim); + const char *delim = ","; + char *token = strtok((char*)opt_arg, delim); ucs_status_t status; - status = parse_mem_type(token, send_mem_type); + status = parse_perf_mem_allocator(token, send_mem_type, send_alloc_name); if (status != UCS_OK) { return status; } @@ -293,9 +326,11 @@ static ucs_status_t parse_mem_type_params(const char *opt_arg, token = strtok(NULL, delim); if (NULL == token) { *recv_mem_type = *send_mem_type; + ucs_strncpy_safe(recv_alloc_name, send_alloc_name, + UCX_PERF_ALLOC_NAME_MAX); return UCS_OK; } else { - return parse_mem_type(token, recv_mem_type); + return parse_perf_mem_allocator(token, recv_mem_type, recv_alloc_name); } } @@ -784,7 +819,9 @@ ucs_status_t parse_test_params(perftest_params_t *params, char opt, } case 'm': return parse_mem_type_params(opt_arg, ¶ms->super.send_mem_type, - ¶ms->super.recv_mem_type); + ¶ms->super.recv_mem_type, + params->super.send_mem_alloc_name, + params->super.recv_mem_alloc_name); case 'a': return parse_accel_device_params(opt_arg, ¶ms->super.send_device, ¶ms->super.recv_device); diff --git a/src/tools/perf/perftest_run.c b/src/tools/perf/perftest_run.c index e3711013ded..fa1be38ad46 100644 --- a/src/tools/perf/perftest_run.c +++ b/src/tools/perf/perftest_run.c @@ -156,8 +156,10 @@ static void print_header(struct perftest_context *ctx) printf("| API: %-60s |\n", test_api_str); printf("| Test: %-60s |\n", test->desc); printf("| Data layout: %-60s |\n", test_data_str); - printf("| Send memory: %-60s |\n", ucs_memory_type_names[ctx->params.super.send_mem_type]); - printf("| Recv memory: %-60s |\n", ucs_memory_type_names[ctx->params.super.recv_mem_type]); + printf("| Send memory: %-60s |\n", + ucx_perf_mem_alloc_name(&ctx->params.super, 1)); + printf("| Recv memory: %-60s |\n", + ucx_perf_mem_alloc_name(&ctx->params.super, 0)); if (ctx->params.super.send_device.mem_type != UCS_MEMORY_TYPE_LAST) { get_accel_device_str(&ctx->params.super.send_device, mem_dev_str, sizeof(mem_dev_str)); printf("| Send device: %-60s |\n", mem_dev_str); diff --git a/src/tools/perf/rocm/rocm_alloc.c b/src/tools/perf/rocm/rocm_alloc.c index a1f1bb3a6c4..baec819dbe8 100644 --- a/src/tools/perf/rocm/rocm_alloc.c +++ b/src/tools/perf/rocm/rocm_alloc.c @@ -142,29 +142,33 @@ static void* ucx_perf_rocm_memset(void *dst, int value, size_t count) return dst; } +static ucx_perf_allocator_t rocm_allocator = { + .name = "rocm", + .default_mem_type = UCS_MEMORY_TYPE_ROCM, + .init = ucx_perf_rocm_init, + .uct_alloc = uct_perf_rocm_alloc, + .uct_free = uct_perf_rocm_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_rocm_memcpy, + .memset = ucx_perf_rocm_memset +}; + +static ucx_perf_allocator_t rocm_managed_allocator = { + .name = "rocm-managed", + .default_mem_type = UCS_MEMORY_TYPE_ROCM_MANAGED, + .init = ucx_perf_rocm_init, + .uct_alloc = uct_perf_rocm_managed_alloc, + .uct_free = uct_perf_rocm_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_rocm_memcpy, + .memset = ucx_perf_rocm_memset +}; + UCS_STATIC_INIT { - static ucx_perf_allocator_t rocm_allocator = { - .mem_type = UCS_MEMORY_TYPE_ROCM, - .init = ucx_perf_rocm_init, - .uct_alloc = uct_perf_rocm_alloc, - .uct_free = uct_perf_rocm_free, - .memcpy = ucx_perf_rocm_memcpy, - .memset = ucx_perf_rocm_memset - }; - static ucx_perf_allocator_t rocm_managed_allocator = { - .mem_type = UCS_MEMORY_TYPE_ROCM_MANAGED, - .init = ucx_perf_rocm_init, - .uct_alloc = uct_perf_rocm_managed_alloc, - .uct_free = uct_perf_rocm_free, - .memcpy = ucx_perf_rocm_memcpy, - .memset = ucx_perf_rocm_memset - }; - - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ROCM] = &rocm_allocator; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ROCM_MANAGED] = &rocm_managed_allocator; + ucx_perf_allocator_register(&rocm_allocator); + ucx_perf_allocator_register(&rocm_managed_allocator); } UCS_STATIC_CLEANUP { - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ROCM] = NULL; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ROCM_MANAGED] = NULL; - + ucx_perf_allocator_unregister(&rocm_managed_allocator); + ucx_perf_allocator_unregister(&rocm_allocator); } diff --git a/src/tools/perf/ze/ze_alloc.c b/src/tools/perf/ze/ze_alloc.c index 8970e617a3e..e66a534107b 100644 --- a/src/tools/perf/ze/ze_alloc.c +++ b/src/tools/perf/ze/ze_alloc.c @@ -285,45 +285,51 @@ static void *ucx_perf_ze_memset(void *dst, int value, size_t count) return dst; } +static ucx_perf_allocator_t ze_host_allocator = { + .name = "ze-host", + .default_mem_type = UCS_MEMORY_TYPE_ZE_HOST, + .init = ucx_perf_ze_init, + .uct_alloc = uct_perf_ze_host_alloc, + .uct_free = uct_perf_ze_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_ze_memcpy, + .memset = ucx_perf_ze_memset +}; + +static ucx_perf_allocator_t ze_device_allocator = { + .name = "ze-device", + .default_mem_type = UCS_MEMORY_TYPE_ZE_DEVICE, + .init = ucx_perf_ze_init, + .uct_alloc = uct_perf_ze_device_alloc, + .uct_free = uct_perf_ze_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_ze_memcpy, + .memset = ucx_perf_ze_memset +}; + +static ucx_perf_allocator_t ze_managed_allocator = { + .name = "ze-managed", + .default_mem_type = UCS_MEMORY_TYPE_ZE_MANAGED, + .init = ucx_perf_ze_init, + .uct_alloc = uct_perf_ze_managed_alloc, + .uct_free = uct_perf_ze_free, + .resolve_mem_type = ucx_perf_allocator_default_resolve_mem_type, + .memcpy = ucx_perf_ze_memcpy, + .memset = ucx_perf_ze_memset +}; + UCS_STATIC_INIT { - static ucx_perf_allocator_t ze_host_allocator = { - .mem_type = UCS_MEMORY_TYPE_ZE_HOST, - .init = ucx_perf_ze_init, - .uct_alloc = uct_perf_ze_host_alloc, - .uct_free = uct_perf_ze_free, - .memcpy = ucx_perf_ze_memcpy, - .memset = ucx_perf_ze_memset - }; - static ucx_perf_allocator_t ze_device_allocator = { - .mem_type = UCS_MEMORY_TYPE_ZE_DEVICE, - .init = ucx_perf_ze_init, - .uct_alloc = uct_perf_ze_device_alloc, - .uct_free = uct_perf_ze_free, - .memcpy = ucx_perf_ze_memcpy, - .memset = ucx_perf_ze_memset - }; - static ucx_perf_allocator_t ze_managed_allocator = { - .mem_type = UCS_MEMORY_TYPE_ZE_MANAGED, - .init = ucx_perf_ze_init, - .uct_alloc = uct_perf_ze_managed_alloc, - .uct_free = uct_perf_ze_free, - .memcpy = ucx_perf_ze_memcpy, - .memset = ucx_perf_ze_memset - }; - - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ZE_HOST] = &ze_host_allocator; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ZE_DEVICE] = - &ze_device_allocator; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ZE_MANAGED] = - &ze_managed_allocator; + ucx_perf_allocator_register(&ze_host_allocator); + ucx_perf_allocator_register(&ze_device_allocator); + ucx_perf_allocator_register(&ze_managed_allocator); } UCS_STATIC_CLEANUP { ucx_perf_ze_destroy_tls_cmdlist(); - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ZE_HOST] = NULL; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ZE_DEVICE] = NULL; - ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_ZE_MANAGED] = NULL; + ucx_perf_allocator_unregister(&ze_managed_allocator); + ucx_perf_allocator_unregister(&ze_device_allocator); + ucx_perf_allocator_unregister(&ze_host_allocator); }