diff --git a/src/ucm/api/ucm.h b/src/ucm/api/ucm.h index 37ccb4592d1..8d04e4860a6 100644 --- a/src/ucm/api/ucm.h +++ b/src/ucm/api/ucm.h @@ -15,6 +15,7 @@ BEGIN_C_DECLS #include #include +#include #include #include @@ -192,11 +193,15 @@ typedef union ucm_event { * Memory type allocation and deallocation event. * If mem_type is @ref UCS_MEMORY_TYPE_LAST, the memory type is unknown, and * further memory type detection is required. + * If sys_dev is @ref UCS_SYS_DEVICE_ID_UNKNOWN and mem_flags is 0, the + * remaining memory attributes are unknown and require detection. */ struct { void *address; size_t size; ucs_memory_type_t mem_type; + ucs_sys_device_t sys_dev; + uint8_t mem_flags; } mem_type; } ucm_event_t; diff --git a/src/ucm/cuda/cudamem.c b/src/ucm/cuda/cudamem.c index f280afdaa5e..a8fd0774d6e 100644 --- a/src/ucm/cuda/cudamem.c +++ b/src/ucm/cuda/cudamem.c @@ -79,7 +79,8 @@ ucm_trace("%s(size_ptr=%p, obj=%p, name=%s) returned dptr=%p", \ __func__, bytes, obj, name, (void*)(*dptr)); \ ucm_cuda_dispatch_mem_type_alloc(*dptr, *size_ptr, \ - UCS_MEMORY_TYPE_CUDA_MANAGED); \ + UCS_MEMORY_TYPE_CUDA_MANAGED, \ + 0); \ } \ ucm_event_leave(); \ return ret; \ @@ -155,20 +156,24 @@ UCM_DEFINE_REPLACE_DLSYM_PTR_FUNC(cudaGetSymbolAddress, cudaError_t, -1, void**, const void*) static void ucm_cuda_dispatch_mem_type_alloc(CUdeviceptr ptr, size_t length, - ucs_memory_type_t mem_type) + ucs_memory_type_t mem_type, + uint8_t mem_flags) { ucm_event_t event; - event.mem_type.address = (void*)ptr; - event.mem_type.size = length; - event.mem_type.mem_type = mem_type; + event.mem_type.address = (void*)ptr; + event.mem_type.size = length; + event.mem_type.mem_type = mem_type; + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = mem_flags; ucm_event_dispatch(UCM_EVENT_MEM_TYPE_ALLOC, &event); } static void ucm_cuda_dispatch_mem_alloc(CUdeviceptr ptr, size_t length) { /* Indicate unknown type and let cuda_md detect attributes. */ - ucm_cuda_dispatch_mem_type_alloc(ptr, length, UCS_MEMORY_TYPE_LAST); + ucm_cuda_dispatch_mem_type_alloc(ptr, length, UCS_MEMORY_TYPE_LAST, + UCS_MEM_FLAG_REGISTRABLE); } static void ucm_cuda_dispatch_mem_free(CUdeviceptr ptr, size_t length, @@ -197,9 +202,11 @@ static void ucm_cuda_dispatch_mem_free(CUdeviceptr ptr, size_t length, } } - event.mem_type.address = (void*)ptr; - event.mem_type.size = length; - event.mem_type.mem_type = mem_type; + event.mem_type.address = (void*)ptr; + event.mem_type.size = length; + event.mem_type.mem_type = mem_type; + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_dispatch(UCM_EVENT_MEM_TYPE_FREE, &event); } @@ -295,7 +302,7 @@ cudaError_t ucm_cudaGetSymbolAddress(void **devPtr, const void *symbol) ucm_trace("%s(symbol=%p) allocated %p", __func__, symbol, *devPtr); ucm_cuda_dispatch_mem_type_alloc((CUdeviceptr)*devPtr, ucm_cuda_get_symbol_size(symbol), - UCS_MEMORY_TYPE_CUDA_MANAGED); + UCS_MEMORY_TYPE_CUDA_MANAGED, 0); } ucm_event_leave(); return ret; @@ -457,9 +464,11 @@ static int ucm_cudamem_scan_regions_cb(void *arg, void *addr, size_t length, ucm_debug("dispatching initial memtype allocation for %p..%p %s", addr, UCS_PTR_BYTE_OFFSET(addr, length), path); - event.mem_type.address = addr; - event.mem_type.size = length; - event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */ + event.mem_type.address = addr; + event.mem_type.size = length; + event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */ + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_enter(); handler->cb(UCM_EVENT_MEM_TYPE_ALLOC, &event, handler->arg); diff --git a/src/ucm/rocm/rocmmem.c b/src/ucm/rocm/rocmmem.c index e213662d806..0470e42958c 100644 --- a/src/ucm/rocm/rocmmem.c +++ b/src/ucm/rocm/rocmmem.c @@ -36,9 +36,11 @@ ucm_dispatch_mem_type_alloc(void *addr, size_t length, ucs_memory_type_t mem_typ { ucm_event_t event; - event.mem_type.address = addr; - event.mem_type.size = length; - event.mem_type.mem_type = mem_type; + event.mem_type.address = addr; + event.mem_type.size = length; + event.mem_type.mem_type = mem_type; + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_dispatch(UCM_EVENT_MEM_TYPE_ALLOC, &event); } @@ -47,9 +49,11 @@ ucm_dispatch_mem_type_free(void *addr, size_t length, ucs_memory_type_t mem_type { ucm_event_t event; - event.mem_type.address = addr; - event.mem_type.size = length; - event.mem_type.mem_type = mem_type; + event.mem_type.address = addr; + event.mem_type.size = length; + event.mem_type.mem_type = mem_type; + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_dispatch(UCM_EVENT_MEM_TYPE_FREE, &event); } @@ -183,9 +187,11 @@ static int ucm_rocm_scan_regions_cb(void *arg, void *addr, size_t length, ucm_debug("dispatching initial memtype allocation for %p..%p %s", addr, UCS_PTR_BYTE_OFFSET(addr, length), path); - event.mem_type.address = addr; - event.mem_type.size = length; - event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */ + event.mem_type.address = addr; + event.mem_type.size = length; + event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */ + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_enter(); handler->cb(UCM_EVENT_MEM_TYPE_ALLOC, &event, handler->arg); diff --git a/src/ucm/ze/zemem.c b/src/ucm/ze/zemem.c index 8aaeedcd61a..f473881fa13 100644 --- a/src/ucm/ze/zemem.c +++ b/src/ucm/ze/zemem.c @@ -45,9 +45,11 @@ ucm_dispatch_mem_type_alloc(void *addr, size_t length, { ucm_event_t event; - event.mem_type.address = addr; - event.mem_type.size = length; - event.mem_type.mem_type = mem_type; + event.mem_type.address = addr; + event.mem_type.size = length; + event.mem_type.mem_type = mem_type; + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_dispatch(UCM_EVENT_MEM_TYPE_ALLOC, &event); } @@ -57,9 +59,11 @@ ucm_dispatch_mem_type_free(void *addr, size_t length, { ucm_event_t event; - event.mem_type.address = addr; - event.mem_type.size = length; - event.mem_type.mem_type = mem_type; + event.mem_type.address = addr; + event.mem_type.size = length; + event.mem_type.mem_type = mem_type; + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_dispatch(UCM_EVENT_MEM_TYPE_FREE, &event); } @@ -234,9 +238,11 @@ static int ucm_zemem_scan_regions_cb(void *arg, void *addr, size_t length, ucm_debug("dispatching initial memtype allocation for %p..%p %s", addr, UCS_PTR_BYTE_OFFSET(addr, length), path); - event.mem_type.address = addr; - event.mem_type.size = length; - event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */ + event.mem_type.address = addr; + event.mem_type.size = length; + event.mem_type.mem_type = UCS_MEMORY_TYPE_LAST; /* unknown memory type */ + event.mem_type.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE; ucm_event_enter(); handler->cb(UCM_EVENT_MEM_TYPE_ALLOC, &event, handler->arg); diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index fbddae774a6..5d5a421243b 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -738,8 +738,8 @@ ucp_memory_detect_internal(ucp_context_h context, const void *address, } else if (ucs_likely(status == UCS_OK)) { if (ucs_unlikely( (mem_info->type == UCS_MEMORY_TYPE_UNKNOWN) || - (mem_info->mem_flags & - UCS_MEM_FLAG_NEEDS_QUERY))) { + ((mem_info->sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) && + (mem_info->mem_flags == 0)))) { ucs_trace_req("address %p length %zu: querying memory attributes", address, length); ucp_memory_detect_slowpath(context, address, length, mem_info); diff --git a/src/ucs/memory/memory_type.h b/src/ucs/memory/memory_type.h index 1d00f3464ed..f380f510a03 100644 --- a/src/ucs/memory/memory_type.h +++ b/src/ucs/memory/memory_type.h @@ -52,8 +52,7 @@ typedef enum ucs_memory_type { typedef enum ucs_mem_flags { - UCS_MEM_FLAG_REGISTRABLE = UCS_BIT(0), /**< Memory is registrable by MDs */ - UCS_MEM_FLAG_NEEDS_QUERY = UCS_BIT(1) /**< MD query is required */ + UCS_MEM_FLAG_REGISTRABLE = UCS_BIT(0) /**< Memory is registrable by MDs */ } ucs_mem_flags_t; diff --git a/src/ucs/memory/memtype_cache.c b/src/ucs/memory/memtype_cache.c index b3a5bbb1450..337825436ea 100644 --- a/src/ucs/memory/memtype_cache.c +++ b/src/ucs/memory/memtype_cache.c @@ -306,7 +306,6 @@ static void ucs_memtype_cache_event_callback(ucm_event_type_t event_type, ucm_event_t *event, void *arg) { ucs_memtype_cache_action_t action; - uint8_t mem_flags; if (event_type & UCM_EVENT_MEM_TYPE_ALLOC) { action = UCS_MEMTYPE_CACHE_ACTION_SET_MEMTYPE; @@ -316,21 +315,17 @@ static void ucs_memtype_cache_event_callback(ucm_event_type_t event_type, return; } - ucs_trace("dispatching mem event %d address %p length %zu mem_type %s", + ucs_trace("dispatching mem event %d address %p length %zu mem_type %s " + "sys_dev %u flags 0x%x", event_type, event->mem_type.address, event->mem_type.size, - ucs_memory_type_names[event->mem_type.mem_type]); - - /* Complete CUDA managed attributes by querying the MD. */ - mem_flags = UCS_MEM_FLAG_REGISTRABLE; - if (event->mem_type.mem_type == UCS_MEMORY_TYPE_CUDA_MANAGED) { - mem_flags |= UCS_MEM_FLAG_NEEDS_QUERY; - } + ucs_memory_type_names[event->mem_type.mem_type], + event->mem_type.sys_dev, event->mem_type.mem_flags); ucs_memtype_cache_update_internal(arg, event->mem_type.address, event->mem_type.size, event->mem_type.mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN, - mem_flags, action); + event->mem_type.sys_dev, + event->mem_type.mem_flags, action); } static void ucs_memtype_cache_purge(ucs_memtype_cache_t *memtype_cache) diff --git a/src/uct/cuda/cuda_copy/cuda_copy_ep.c b/src/uct/cuda/cuda_copy/cuda_copy_ep.c index 3947907972f..9e17ba8200c 100644 --- a/src/uct/cuda/cuda_copy/cuda_copy_ep.c +++ b/src/uct/cuda/cuda_copy/cuda_copy_ep.c @@ -87,8 +87,7 @@ uct_cuda_copy_get_mem_type(uct_md_h md, const void *address, size_t length, if (ucs_unlikely((status == UCS_ERR_UNSUPPORTED) || (mem_info.type == UCS_MEMORY_TYPE_UNKNOWN) || - (mem_info.mem_flags & - UCS_MEM_FLAG_NEEDS_QUERY))) { + (mem_info.sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN))) { mem_attr.field_mask = UCT_MD_MEM_ATTR_V2_FIELD_MEM_TYPE | UCT_MD_MEM_ATTR_V2_FIELD_SYS_DEV; diff --git a/test/apps/test_cuda_get_symbol_address.cu b/test/apps/test_cuda_get_symbol_address.cu index 09536dbaaaa..d9430836466 100644 --- a/test/apps/test_cuda_get_symbol_address.cu +++ b/test/apps/test_cuda_get_symbol_address.cu @@ -48,6 +48,20 @@ static void event_cb(ucm_event_type_t event_type, ucm_event_t *event, void *arg) ucs_memory_type_names[UCS_MEMORY_TYPE_CUDA_MANAGED]); ++ctx->num_errors; } + + if ((event_type == UCM_EVENT_MEM_TYPE_ALLOC) && + (event->mem_type.sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN)) { + printf("unexpected symbol system device %u, expected unknown\n", + event->mem_type.sys_dev); + ++ctx->num_errors; + } + + if ((event_type == UCM_EVENT_MEM_TYPE_ALLOC) && + (event->mem_type.mem_flags != 0)) { + printf("unexpected symbol memory flags 0x%x, expected 0\n", + event->mem_type.mem_flags); + ++ctx->num_errors; + } } static int check_symbol_memory_type(ucp_context_h context, void *address) diff --git a/test/gtest/ucs/test_memtype_cache.cc b/test/gtest/ucs/test_memtype_cache.cc index 0f9822b3f18..475bb2ba717 100644 --- a/test/gtest/ucs/test_memtype_cache.cc +++ b/test/gtest/ucs/test_memtype_cache.cc @@ -479,6 +479,30 @@ UCS_TEST_P(test_memtype_cache, diff_mem_types_diff_bufs_keep_mem) { test_memtype_cache_alloc_diff_mem_types(true, false); } +UCS_TEST_P(test_memtype_cache, event_attributes) { + const size_t size = ucs_get_page_size(); + void *address = reinterpret_cast(0xdead0000ul); + const ucs_sys_device_t dev = 0; + ucm_event_t event; + ucs_memory_info_t mem_info; + + event.mem_type.address = address; + event.mem_type.size = size; + event.mem_type.mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED; + event.mem_type.sys_dev = dev; + event.mem_type.mem_flags = 0; + ucm_event_dispatch(UCM_EVENT_MEM_TYPE_ALLOC, &event); + + ASSERT_UCS_OK(ucs_memtype_cache_lookup(address, size, &mem_info)); + EXPECT_EQ(event.mem_type.mem_type, mem_info.type); + EXPECT_EQ(dev, mem_info.sys_dev); + EXPECT_EQ(0, mem_info.mem_flags); + + ucm_event_dispatch(UCM_EVENT_MEM_TYPE_FREE, &event); + EXPECT_EQ(UCS_ERR_NO_ELEM, + ucs_memtype_cache_lookup(address, size, &mem_info)); +} + INSTANTIATE_TEST_SUITE_P(mem_type, test_memtype_cache, ::testing::ValuesIn(mem_buffer::supported_mem_types()));