diff --git a/src/uct/cuda/Makefile.am b/src/uct/cuda/Makefile.am index bd11d39fd09..e171da37e44 100644 --- a/src/uct/cuda/Makefile.am +++ b/src/uct/cuda/Makefile.am @@ -24,6 +24,7 @@ noinst_HEADERS = \ base/cuda_ctx.h \ base/cuda_ctx.inl \ base/cuda_md.h \ + base/cuda_mem.h \ base/cuda_iface.h \ base/cuda_nvml.h \ base/cuda_util.h \ @@ -41,6 +42,7 @@ libuct_cuda_la_SOURCES = \ base/cuda_ctx.c \ base/cuda_iface.c \ base/cuda_md.c \ + base/cuda_mem.c \ base/cuda_nvml.c \ base/cuda_util.c \ cuda_copy/cuda_copy_md.c \ diff --git a/src/uct/cuda/base/cuda_mem.c b/src/uct/cuda/base/cuda_mem.c new file mode 100644 index 00000000000..9e660fb449c --- /dev/null +++ b/src/uct/cuda/base/cuda_mem.c @@ -0,0 +1,181 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#ifdef HAVE_CONFIG_H +#include "config.h" +#endif + +#include "cuda_mem.h" + +#include "cuda_util.h" + +#include +#include + +static ucs_status_t uct_cuda_mem_alloc_fabric(ucs_log_level_t log_level, + CUdevice cu_device, size_t length, + size_t *granularity, + uct_cuda_mem_t *mem) +{ +#if HAVE_CUDA_FABRIC + CUmemAllocationProp prop = {}; + CUmemAccessDesc access_desc = {}; + ucs_status_t status; + uint64_t allowed_types; + + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = cu_device; + prop.allocFlags.gpuDirectRDMACapable = 1; + + if (*granularity == SIZE_MAX) { + status = UCT_CUDADRV_FUNC( + cuMemGetAllocationGranularity(granularity, &prop, + CU_MEM_ALLOC_GRANULARITY_MINIMUM), + log_level); + if (status != UCS_OK) { + return status; + } + } + + mem->length = ucs_align_up(length, *granularity); + + status = UCT_CUDADRV_FUNC(cuMemCreate(&mem->generic_handle, mem->length, + &prop, 0), + log_level); + if (status != UCS_OK) { + return UCS_ERR_NO_MEMORY; + } + + status = UCT_CUDADRV_FUNC(cuMemAddressReserve(&mem->ptr, mem->length, + *granularity, 0, 0), + log_level); + if (status != UCS_OK) { + goto err_mem_release; + } + + status = UCT_CUDADRV_FUNC(cuMemMap(mem->ptr, mem->length, 0, + mem->generic_handle, 0), + log_level); + if (status != UCS_OK) { + goto err_address_free; + } + + access_desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + access_desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access_desc.location.id = cu_device; + + status = UCT_CUDADRV_FUNC(cuMemSetAccess(mem->ptr, mem->length, + &access_desc, 1), + log_level); + if (status != UCS_OK) { + goto err_mem_unmap; + } + + status = UCT_CUDADRV_FUNC( + cuPointerGetAttribute(&allowed_types, + CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES, + mem->ptr), + log_level); + if (status != UCS_OK) { + goto err_mem_unmap; + } else if (!(allowed_types & CU_MEM_HANDLE_TYPE_FABRIC)) { + ucs_log(log_level, + "allocated memory at %p of size %ld does not have fabric " + "property", + (void*)mem->ptr, mem->length); + goto err_mem_unmap; + } + + mem->is_vmm = 1; + + ucs_trace("allocated vmm fabric memory at %p of size %ld", (void*)mem->ptr, + mem->length); + return UCS_OK; + +err_mem_unmap: + (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemUnmap(mem->ptr, mem->length)); +err_address_free: + (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemAddressFree(mem->ptr, mem->length)); +err_mem_release: + (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemRelease(mem->generic_handle)); +#endif + return UCS_ERR_NO_MEMORY; +} + +static void uct_cuda_mem_release_fabric(uct_cuda_mem_t *mem) +{ +#if HAVE_CUDA_FABRIC + (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemUnmap(mem->ptr, mem->length)); + (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemAddressFree(mem->ptr, mem->length)); + (void)UCT_CUDADRV_FUNC_LOG_WARN(cuMemRelease(mem->generic_handle)); +#endif +} + +ucs_status_t +uct_cuda_mem_alloc(ucs_log_level_t log_level, ucs_memory_type_t mem_type, + ucs_ternary_auto_value_t enable_fabric, CUdevice cu_device, + size_t length, size_t *granularity_p, uct_cuda_mem_t *mem) +{ + CUdeviceptr ptr; + ucs_status_t status; + ucs_log_level_t log_level_fabric; + + if ((mem_type != UCS_MEMORY_TYPE_CUDA) && + (mem_type != UCS_MEMORY_TYPE_CUDA_MANAGED)) { + ucs_log(log_level, + "unsupported memory type %s, supported types: %s, %s", + ucs_memory_type_names[mem_type], + ucs_memory_type_names[UCS_MEMORY_TYPE_CUDA], + ucs_memory_type_names[UCS_MEMORY_TYPE_CUDA_MANAGED]); + return UCS_ERR_INVALID_PARAM; + } + + if (mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED) { + status = UCT_CUDADRV_FUNC(cuMemAllocManaged(&ptr, length, + CU_MEM_ATTACH_GLOBAL), + log_level); + if (status != UCS_OK) { + return status; + } + + goto set_non_vmm; + } + + if (enable_fabric != UCS_NO) { + log_level_fabric = (enable_fabric == UCS_TRY) ? UCS_LOG_LEVEL_DEBUG : + log_level; + status = uct_cuda_mem_alloc_fabric(log_level_fabric, cu_device, length, + granularity_p, mem); + if ((status == UCS_OK) || (enable_fabric == UCS_YES)) { + return status; + } + } + + status = UCT_CUDADRV_FUNC(cuMemAlloc(&ptr, length), log_level); + if (status != UCS_OK) { + return UCS_ERR_NO_MEMORY; + } + +set_non_vmm: + mem->ptr = ptr; + mem->length = length; + mem->is_vmm = 0; + + return UCS_OK; +} + +void uct_cuda_mem_free(uct_cuda_mem_t *mem) +{ + if (mem->is_vmm) { + uct_cuda_mem_release_fabric(mem); + } else { + /* TODO: The more correct behavior is to issue a warning, provided that + * the context still exists for which memory has been allocated */ + (void)UCT_CUDADRV_FUNC(cuMemFree(mem->ptr), UCS_LOG_LEVEL_DIAG); + } +} diff --git a/src/uct/cuda/base/cuda_mem.h b/src/uct/cuda/base/cuda_mem.h new file mode 100644 index 00000000000..2872a52dd2d --- /dev/null +++ b/src/uct/cuda/base/cuda_mem.h @@ -0,0 +1,58 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#ifndef UCT_CUDA_MEM_H +#define UCT_CUDA_MEM_H + +#include +#include +#include +#include + +#include + +#include +#include + +typedef struct { + CUdeviceptr ptr; + size_t length; + uint8_t is_vmm; +#if HAVE_CUDA_FABRIC + CUmemGenericAllocationHandle generic_handle; +#endif +} uct_cuda_mem_t; + + +/** + * Allocate CUDA device memory, optionally using fabric VMM allocation. + * + * @param [in] log_level Log level for CUDA driver API failures + * @param [in] mem_type Memory type to allocate + * @param [in] enable_fabric Controls fabric VMM allocation + * @param [in] cu_device CUDA device to allocate memory on + * @param [in] length The minimal size to allocate + * @param [in,out] granularity_p Fabric VMM allocation granularity. If the + * pointed-to value is SIZE_MAX, the granularity + * is queried and written back. + * @param [out] mem Filled with information about the allocated + * memory + */ +ucs_status_t +uct_cuda_mem_alloc(ucs_log_level_t log_level, ucs_memory_type_t mem_type, + ucs_ternary_auto_value_t enable_fabric, CUdevice cu_device, + size_t length, size_t *granularity_p, uct_cuda_mem_t *mem); + + +/** + * Release the memory allocated by @ref uct_cuda_mem_alloc. + * + * @param [in] mem Description of allocated memory, as returned from + * @ref uct_cuda_mem_alloc + */ +void uct_cuda_mem_free(uct_cuda_mem_t *mem); + +#endif /* UCT_CUDA_MEM_H */ diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.c b/src/uct/cuda/cuda_copy/cuda_copy_md.c index bada99db7f7..1dbc40e6c19 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.c @@ -9,8 +9,6 @@ #include "cuda_copy_md.h" -#include -#include #include #include #include @@ -20,17 +18,20 @@ #include #include #include +#include #include #include + #include #if CUDA_VERSION >= 11070 #include #endif +#include +#include #define UCT_CUDA_DEV_NAME_MAX_LEN 64 - static const char *uct_cuda_pref_loc[] = { [UCT_CUDA_PREF_LOC_CPU] = "cpu", [UCT_CUDA_PREF_LOC_GPU] = "gpu", @@ -191,101 +192,6 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_copy_mem_dereg, return UCS_OK; } -static ucs_status_t -uct_cuda_copy_mem_alloc_fabric(ucs_log_level_t log_level, CUdevice cu_device, - size_t length, size_t *granularity, - uct_cuda_copy_alloc_handle_t *alloc_handle) -{ -#if HAVE_CUDA_FABRIC - CUmemAllocationProp prop = {}; - CUmemAccessDesc access_desc = {}; - ucs_status_t status; - uint64_t allowed_types; - - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = cu_device; - prop.allocFlags.gpuDirectRDMACapable = 1; - - if (*granularity == SIZE_MAX) { - status = UCT_CUDADRV_FUNC( - cuMemGetAllocationGranularity(granularity, &prop, - CU_MEM_ALLOC_GRANULARITY_MINIMUM), - log_level); - if (status != UCS_OK) { - return status; - } - } - - alloc_handle->length = ucs_align_up(length, *granularity); - - status = UCT_CUDADRV_FUNC(cuMemCreate(&alloc_handle->generic_handle, - alloc_handle->length, &prop, 0), - log_level); - if (status != UCS_OK) { - return UCS_ERR_NO_MEMORY; - } - - status = UCT_CUDADRV_FUNC(cuMemAddressReserve(&alloc_handle->ptr, - alloc_handle->length, - *granularity, 0, 0), - log_level); - if (status != UCS_OK) { - goto err_mem_release; - } - - status = UCT_CUDADRV_FUNC(cuMemMap(alloc_handle->ptr, alloc_handle->length, - 0, alloc_handle->generic_handle, 0), - log_level); - if (status != UCS_OK) { - goto err_address_free; - } - - access_desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - access_desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - access_desc.location.id = cu_device; - - status = UCT_CUDADRV_FUNC(cuMemSetAccess( - alloc_handle->ptr, alloc_handle->length, &access_desc, 1), - log_level); - if (status != UCS_OK) { - goto err_mem_unmap; - } - - status = UCT_CUDADRV_FUNC( - cuPointerGetAttribute(&allowed_types, - CU_POINTER_ATTRIBUTE_ALLOWED_HANDLE_TYPES, - alloc_handle->ptr), - log_level); - if (status != UCS_OK) { - goto err_mem_unmap; - } else if (!(allowed_types & CU_MEM_HANDLE_TYPE_FABRIC)) { - ucs_log(log_level, - "allocated memory at %p of size %ld does not have fabric " - "property", - (void*)alloc_handle->ptr, alloc_handle->length); - goto err_mem_unmap; - } - - alloc_handle->is_vmm = 1; - - ucs_trace("allocated vmm fabric memory at %p of size %ld", - (void*)alloc_handle->ptr, alloc_handle->length); - return UCS_OK; - -err_mem_unmap: - UCT_CUDADRV_FUNC_LOG_DEBUG( - cuMemUnmap(alloc_handle->ptr, alloc_handle->length)); -err_address_free: - UCT_CUDADRV_FUNC_LOG_DEBUG( - cuMemAddressFree(alloc_handle->ptr, alloc_handle->length)); -err_mem_release: - UCT_CUDADRV_FUNC_LOG_DEBUG(cuMemRelease(alloc_handle->generic_handle)); -#endif - return UCS_ERR_NO_MEMORY; -} - typedef CUresult (*uct_cuda_cuCtxSetFlags_t)(unsigned); ucs_status_t uct_cuda_copy_set_ctx_sync_memops(ucs_log_level_t log_level) @@ -359,7 +265,7 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p, uct_cuda_copy_md_t *md = ucs_derived_of(uct_md, uct_cuda_copy_md_t); uct_md_mem_attr_v2_t mem_attr; ucs_status_t status; - uct_cuda_copy_alloc_handle_t *alloc_handle; + uct_cuda_mem_t *cuda_mem; ucs_log_level_t log_level; CUdevice avail_cuda_device, cuda_device; @@ -371,11 +277,9 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p, log_level = (flags & UCT_MD_MEM_FLAG_HIDE_ERRORS) ? UCS_LOG_LEVEL_DEBUG : UCS_LOG_LEVEL_ERROR; - alloc_handle = ucs_malloc(sizeof(*alloc_handle), - "uct_cuda_copy_alloc_handle_t"); - if (NULL == alloc_handle) { - ucs_log(log_level, - "failed to allocate memory for uct_cuda_copy_alloc_handle_t"); + cuda_mem = ucs_malloc(sizeof(*cuda_mem), "uct_cuda_mem_t"); + if (NULL == cuda_mem) { + ucs_log(log_level, "failed to allocate memory for uct_cuda_mem_t"); return UCS_ERR_NO_MEMORY; } @@ -383,19 +287,19 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p, sys_dev, &cuda_device, &avail_cuda_device, log_level); if (status != UCS_OK) { - ucs_free(alloc_handle); + ucs_free(cuda_mem); return UCS_ERR_NO_DEVICE; } status = uct_cuda_mem_alloc(log_level, mem_type, md->config.enable_fabric, avail_cuda_device, *length_p, &md->granularity, - alloc_handle); + cuda_mem); if (status != UCS_OK) { - ucs_free(alloc_handle); + ucs_free(cuda_mem); goto out; } - uct_cuda_copy_sync_memops(alloc_handle->ptr, alloc_handle->is_vmm); + uct_cuda_copy_sync_memops(cuda_mem->ptr, cuda_mem->is_vmm); /* Cache memory flags as part of uct_cuda_copy_md_mem_query() before * restoring the CUDA context. @@ -406,16 +310,16 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p, UCT_MD_MEM_ATTR_V2_FIELD_ALLOC_LENGTH | UCT_MD_MEM_ATTR_V2_FIELD_MEM_FLAGS; - status = uct_cuda_copy_md_mem_query(uct_md, (void*)alloc_handle->ptr, - alloc_handle->length, &mem_attr); + status = uct_cuda_copy_md_mem_query(uct_md, (void*)cuda_mem->ptr, + cuda_mem->length, &mem_attr); if (status != UCS_OK) { - (void)uct_md_mem_free(uct_md, alloc_handle); + (void)uct_md_mem_free(uct_md, cuda_mem); goto out; } - *memh_p = alloc_handle; - *address_p = (void*)alloc_handle->ptr; - *length_p = alloc_handle->length; + *memh_p = cuda_mem; + *address_p = (void*)cuda_mem->ptr; + *length_p = cuda_mem->length; out: if (cuda_device != avail_cuda_device) { @@ -425,31 +329,6 @@ uct_cuda_copy_mem_alloc(uct_md_h uct_md, size_t *length_p, void **address_p, return status; } -static ucs_status_t -uct_cuda_copy_mem_release_fabric(uct_cuda_copy_alloc_handle_t *alloc_handle) -{ -#if HAVE_CUDA_FABRIC - ucs_status_t status; - - status = UCT_CUDADRV_FUNC_LOG_ERR( - cuMemRelease(alloc_handle->generic_handle)); - if (status != UCS_OK) { - return status; - } - - status = UCT_CUDADRV_FUNC_LOG_ERR( - cuMemUnmap(alloc_handle->ptr, alloc_handle->length)); - if (status != UCS_OK) { - return status; - } - - return UCT_CUDADRV_FUNC_LOG_ERR( - cuMemAddressFree(alloc_handle->ptr, alloc_handle->length)); -#else - return UCS_ERR_UNSUPPORTED; -#endif -} - static int uct_cuda_copy_detect_vmm(const void *address, ucs_memory_type_t *vmm_mem_type, CUdevice *cuda_device, @@ -510,11 +389,10 @@ static int uct_cuda_copy_detect_vmm(const void *address, static ucs_status_t uct_cuda_copy_mem_free(uct_md_h md, uct_mem_h memh) { - uct_cuda_copy_alloc_handle_t *alloc_handle = (uct_cuda_copy_alloc_handle_t*) - memh; + uct_cuda_mem_t *cuda_mem = (uct_cuda_mem_t*)memh; - uct_cuda_mem_free(alloc_handle); - ucs_free(alloc_handle); + uct_cuda_mem_free(cuda_mem); + ucs_free(cuda_mem); return UCS_OK; } @@ -1200,70 +1078,3 @@ uct_component_t uct_cuda_copy_component = { .md_vfs_init = (uct_component_md_vfs_init_func_t)ucs_empty_function }; UCT_COMPONENT_REGISTER(&uct_cuda_copy_component); - -ucs_status_t -uct_cuda_mem_alloc(ucs_log_level_t log_level, ucs_memory_type_t mem_type, - ucs_ternary_auto_value_t enable_fabric, CUdevice cu_device, - size_t length, size_t *granularity_p, - uct_cuda_copy_alloc_handle_t *alloc_handle) -{ - CUdeviceptr ptr; - ucs_status_t status; - ucs_log_level_t log_level_fabric; - - if ((mem_type != UCS_MEMORY_TYPE_CUDA) && - (mem_type != UCS_MEMORY_TYPE_CUDA_MANAGED)) { - ucs_log(log_level, - "unsupported memory type %s, supported types: %s, %s", - ucs_memory_type_names[mem_type], - ucs_memory_type_names[UCS_MEMORY_TYPE_CUDA], - ucs_memory_type_names[UCS_MEMORY_TYPE_CUDA_MANAGED]); - return UCS_ERR_INVALID_PARAM; - } - - if (mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED) { - status = UCT_CUDADRV_FUNC(cuMemAllocManaged(&ptr, length, - CU_MEM_ATTACH_GLOBAL), - log_level); - if (status != UCS_OK) { - return status; - } - - goto set_non_vmm; - } - - if (enable_fabric != UCS_NO) { - log_level_fabric = (enable_fabric == UCS_TRY) ? UCS_LOG_LEVEL_DEBUG : - log_level; - status = uct_cuda_copy_mem_alloc_fabric(log_level_fabric, cu_device, - length, granularity_p, - alloc_handle); - if ((status == UCS_OK) || (enable_fabric == UCS_YES)) { - return status; - } - } - - status = UCT_CUDADRV_FUNC(cuMemAlloc(&ptr, length), log_level); - if (status != UCS_OK) { - return UCS_ERR_NO_MEMORY; - } - -set_non_vmm: - alloc_handle->ptr = ptr; - alloc_handle->length = length; - alloc_handle->is_vmm = 0; - - return UCS_OK; -} - -void uct_cuda_mem_free(uct_cuda_copy_alloc_handle_t *alloc_handle) -{ - if (alloc_handle->is_vmm) { - (void)uct_cuda_copy_mem_release_fabric(alloc_handle); - } else { - /* TODO: The more correct behavior is to issue a warning, provided that - * the context still exists for which memory has been allocated */ - (void)UCT_CUDADRV_FUNC(cuMemFree(alloc_handle->ptr), - UCS_LOG_LEVEL_DIAG); - } -} diff --git a/src/uct/cuda/cuda_copy/cuda_copy_md.h b/src/uct/cuda/cuda_copy/cuda_copy_md.h index 2b42802a604..6527f915980 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_md.h +++ b/src/uct/cuda/cuda_copy/cuda_copy_md.h @@ -53,18 +53,6 @@ typedef struct uct_cuda_copy_md_config { int retain_primary_ctx; } uct_cuda_copy_md_config_t; -/** - * copy alloc handle. - */ -typedef struct uct_cuda_copy_alloc_handle { - CUdeviceptr ptr; - size_t length; - uint8_t is_vmm; -#if HAVE_CUDA_FABRIC - CUmemGenericAllocationHandle generic_handle; -#endif -} uct_cuda_copy_alloc_handle_t; - typedef struct { int fd; @@ -105,36 +93,6 @@ uct_cuda_copy_md_dmabuf_t uct_cuda_copy_md_get_dmabuf(const void *address, ucs_sys_device_t sys_dev); -/** - * Allocate CUDA device memory, optionally using fabric VMM allocation. - * - * @param [in] log_level Log level for CUDA driver API failures - * @param [in] mem_type Memory type to allocate - * @param [in] enable_fabric Controls fabric VMM allocation - * @param [in] cu_device CUDA device to allocate memory on - * @param [in] length The minimal size to allocate - * @param [in,out] granularity_p Fabric VMM allocation granularity. If the - * pointed-to value is SIZE_MAX, the granularity - * is queried and written back. - * @param [out] alloc_handle Filled with information about the allocated - * memory - */ -ucs_status_t -uct_cuda_mem_alloc(ucs_log_level_t log_level, ucs_memory_type_t mem_type, - ucs_ternary_auto_value_t enable_fabric, CUdevice cu_device, - size_t length, size_t *granularity_p, - uct_cuda_copy_alloc_handle_t *alloc_handle); - - -/** - * Release the memory allocated by @ref uct_cuda_mem_alloc. - * - * @param [in] alloc_handle Description of allocated memory, as returned from - * @ref uct_cuda_mem_alloc - */ -void uct_cuda_mem_free(uct_cuda_copy_alloc_handle_t *alloc_handle); - - /** * Set the context flag to synchronize DMA operations. * diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c index cbe12f94706..bb157726478 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.c +++ b/src/uct/ib/mlx5/gdaki/gdaki.c @@ -22,6 +22,7 @@ #include #include #include +#include #include @@ -92,9 +93,8 @@ ucs_config_field_t uct_rc_gdaki_iface_config_table[] = { {NULL} }; -static ucs_status_t -uct_gdaki_cuda_mem_alloc(size_t size, size_t *granularity_p, - uct_cuda_copy_alloc_handle_t *alloc_handle) +static ucs_status_t uct_gdaki_cuda_mem_alloc(size_t size, size_t *granularity_p, + uct_cuda_mem_t *cuda_mem) { CUdevice cu_device; ucs_status_t status; @@ -106,26 +106,26 @@ uct_gdaki_cuda_mem_alloc(size_t size, size_t *granularity_p, return uct_cuda_mem_alloc(UCS_LOG_LEVEL_ERROR, UCS_MEMORY_TYPE_CUDA, UCS_TRY, cu_device, size, granularity_p, - alloc_handle); + cuda_mem); } -static ucs_status_t -uct_rc_gdaki_alloc(size_t size, size_t align, size_t *granularity_p, - uct_cuda_copy_alloc_handle_t *alloc_handle, void **buf_p) +static ucs_status_t uct_rc_gdaki_alloc(size_t size, size_t align, + size_t *granularity_p, + uct_cuda_mem_t *cuda_mem, void **buf_p) { unsigned int flag = 1; ucs_status_t status; void *buf; status = uct_gdaki_cuda_mem_alloc(size + align - 1, granularity_p, - alloc_handle); + cuda_mem); if (status != UCS_OK) { return status; } - buf = (void*)ucs_align_up_pow2_ptr(alloc_handle->ptr, align); + buf = (void*)ucs_align_up_pow2_ptr(cuda_mem->ptr, align); - if (!alloc_handle->is_vmm) { + if (!cuda_mem->is_vmm) { status = UCT_CUDADRV_FUNC_LOG_ERR( cuPointerSetAttribute(&flag, CU_POINTER_ATTRIBUTE_SYNC_MEMOPS, (CUdeviceptr)buf)); @@ -138,7 +138,7 @@ uct_rc_gdaki_alloc(size_t size, size_t align, size_t *granularity_p, return UCS_OK; err: - uct_cuda_mem_free(alloc_handle); + uct_cuda_mem_free(cuda_mem); return status; } @@ -166,17 +166,17 @@ static int uct_gdaki_check_umem_dmabuf(const uct_ib_md_t *md) #if HAVE_DECL_MLX5DV_UMEM_MASK_DMABUF size_t granularity = SIZE_MAX; struct mlx5dv_devx_umem_in umem_in = {}; - uct_cuda_copy_alloc_handle_t alloc_handle; + uct_cuda_mem_t cuda_mem; ucs_status_t status; struct mlx5dv_devx_umem *umem; uct_cuda_copy_md_dmabuf_t dmabuf; - status = uct_gdaki_cuda_mem_alloc(1, &granularity, &alloc_handle); + status = uct_gdaki_cuda_mem_alloc(1, &granularity, &cuda_mem); if (status != UCS_OK) { goto out; } - dmabuf = uct_cuda_copy_md_get_dmabuf((void*)alloc_handle.ptr, 1, + dmabuf = uct_cuda_copy_md_get_dmabuf((void*)cuda_mem.ptr, 1, UCS_SYS_DEVICE_ID_UNKNOWN); if (dmabuf.fd == UCT_DMABUF_FD_INVALID) { goto out_free; @@ -200,7 +200,7 @@ static int uct_gdaki_check_umem_dmabuf(const uct_ib_md_t *md) out_close: ucs_close_fd(&dmabuf.fd); out_free: - uct_cuda_mem_free(&alloc_handle); + uct_cuda_mem_free(&cuda_mem); out: #endif return ret; diff --git a/src/uct/ib/mlx5/gdaki/gdaki.h b/src/uct/ib/mlx5/gdaki/gdaki.h index 0944f2576d5..05a8cf43740 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.h +++ b/src/uct/ib/mlx5/gdaki/gdaki.h @@ -8,7 +8,7 @@ #include #include -#include +#include #include #include @@ -30,9 +30,9 @@ typedef struct { } uct_rc_gdaki_channel_block_t; typedef struct { - void *gpu_mem; - uct_cuda_copy_alloc_handle_t gpu_raw; - struct mlx5dv_devx_umem *umem; + void *gpu_mem; + uct_cuda_mem_t gpu_raw; + struct mlx5dv_devx_umem *umem; } uct_rc_gdaki_channel_block_mem_t; typedef struct uct_rc_gdaki_iface {