diff --git a/source/common/unified_malloc_framework/CMakeLists.txt b/source/common/unified_malloc_framework/CMakeLists.txt index 2d3003e79f..176b7ff016 100644 --- a/source/common/unified_malloc_framework/CMakeLists.txt +++ b/source/common/unified_malloc_framework/CMakeLists.txt @@ -8,6 +8,7 @@ set(UMF_SOURCES src/memory_provider.cpp src/memory_tracker.cpp src/memory_provider_get_last_failed.cpp + src/ipc.c ) if(UMF_BUILD_SHARED_LIBRARY) diff --git a/source/common/unified_malloc_framework/include/umf/ipc.h b/source/common/unified_malloc_framework/include/umf/ipc.h new file mode 100644 index 0000000000..9187306e3a --- /dev/null +++ b/source/common/unified_malloc_framework/include/umf/ipc.h @@ -0,0 +1,56 @@ +/* + * + * Copyright (C) 2023 Intel Corporation + * + * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions. + * See LICENSE.TXT + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + * + */ +#ifndef UMF_IPC_H +#define UMF_IPC_H 1 + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct umf_ipc_data_t *umf_ipc_handle_t; + +/// +/// \brief Creates an IPC handle for the specified UMF allocation. +/// \param ptr pointer to the allocated memory. +/// \param ipcHandle [out] returned IPC handle. +/// \param size [out] size of IPC handle in bytes. +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t umfGetIPCHandle(const void *ptr, umf_ipc_handle_t *ipcHandle, + size_t *size); + +/// +/// \brief Release IPC handle retrieved by umfGetIPCHandle. +/// \param ipcHandle IPC handle. +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t umfPutIPCHandle(umf_ipc_handle_t ipcHandle); + +/// +/// \brief Open IPC handle retrieved by umfGetIPCHandle. +/// \param hPool [in] Pool handle where to open the the IPC handle. +/// \param ipcHandle [in] IPC handle. +/// \param ptr [out] pointer to the memory in the current process. +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t umfOpenIPCHandle(umf_memory_pool_handle_t hPool, + umf_ipc_handle_t ipcHandle, void **ptr); + +/// +/// \brief Close IPC handle. +/// \param ptr [in] pointer to the memory. +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t umfCloseIPCHandle(void *ptr); + +#ifdef __cplusplus +} +#endif + +#endif /* UMF_IPC_H */ diff --git a/source/common/unified_malloc_framework/include/umf/memory_provider.h b/source/common/unified_malloc_framework/include/umf/memory_provider.h index 0718c8efb1..d02553137c 100644 --- a/source/common/unified_malloc_framework/include/umf/memory_provider.h +++ b/source/common/unified_malloc_framework/include/umf/memory_provider.h @@ -145,6 +145,46 @@ enum umf_result_t umfMemoryProviderPurgeForce(umf_memory_provider_handle_t hProvider, void *ptr, size_t size); +/// +/// @brief return the size of ipc data structure created by the memory provider +/// \param hProvider [in] handle to the memory provider +/// \param size [out] size of the ipc data structure created by the memory provider +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t +umfMemoryProviderGetIPCHandleSize(umf_memory_provider_handle_t hProvider, + size_t *size); + +/// +/// @brief return IPC handle for the specified base address. +/// \param hProvider [in] handle to the memory provider. +/// \param ptr [in] beginning of the virtual memory range returned by +/// umfMemoryProviderAlloc function. +/// \param size [in] size of the memory address range. +/// \param ipcData [out] IPC handle data. +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t +umfMemoryProviderGetIPCHandle(umf_memory_provider_handle_t hProvider, + const void *ptr, size_t size, void *ipcData); + +/// +/// @brief open IPC handle and return pointer in the current process. +/// \param hProvider [in] handle to the memory provider. +/// \param ipcData [in] IPC handle data. +/// \param ptr [out] returned pointer. +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t +umfMemoryProviderOpenIPCHandle(umf_memory_provider_handle_t hProvider, + void *ipcData, void **ptr); + +/// +/// @brief close IPC handle open with umfMemoryProviderOpenIPCHandle function. +/// \param hProvider [in] handle to the memory provider. +/// \param ptr [in] pointer returned by umfMemoryProviderOpenIPCHandle function. +/// \return UMF_RESULT_SUCCESS on success or appropriate error code on failure. +enum umf_result_t +umfMemoryProviderCloseIPCHandle(umf_memory_provider_handle_t hProvider, + void *ptr); + /// /// \brief Retrieve name of a given memory provider. /// \param hProvider handle to the memory provider diff --git a/source/common/unified_malloc_framework/include/umf/memory_provider_ops.h b/source/common/unified_malloc_framework/include/umf/memory_provider_ops.h index 9a802a2102..7635ac9e1c 100644 --- a/source/common/unified_malloc_framework/include/umf/memory_provider_ops.h +++ b/source/common/unified_malloc_framework/include/umf/memory_provider_ops.h @@ -52,6 +52,12 @@ struct umf_memory_provider_ops_t { enum umf_result_t (*purge_force)(void *provider, void *ptr, size_t size); const char *(*get_name)(void *provider); bool (*supports_device)(const char *name); + enum umf_result_t (*get_ipc_handle_size)(void *provider, size_t *size); + enum umf_result_t (*get_ipc_handle)(void *provider, const void *ptr, + size_t size, void *ipcData); + enum umf_result_t (*open_ipc_handle)(void *provider, void *ipcData, + void **ptr); + enum umf_result_t (*close_ipc_handle)(void *provider, void *ptr); }; #ifdef __cplusplus diff --git a/source/common/unified_malloc_framework/src/ipc.c b/source/common/unified_malloc_framework/src/ipc.c new file mode 100644 index 0000000000..4973ce5a55 --- /dev/null +++ b/source/common/unified_malloc_framework/src/ipc.c @@ -0,0 +1,98 @@ +/* + * + * Copyright (C) 2023 Intel Corporation + * + * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions. + * See LICENSE.TXT + * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + * + */ + +#include "umf/ipc.h" + +#include "memory_pool_internal.h" +#include "memory_tracker.h" + +#include +#include + +struct umf_ipc_data_t { + uint32_t size; + uint64_t offset; + char providerData[]; +}; + +enum umf_result_t +umfGetIPCHandle(const void *ptr, umf_ipc_handle_t *umfIPCHandle, size_t *size) { + size_t ipcHandleSize = 0; + struct umf_alloc_info_t allocInfo; + enum umf_result_t res = + umfMemoryTrackerGetAllocInfo(umfMemoryTrackerGet(), ptr, &allocInfo); + if (res != UMF_RESULT_SUCCESS) { + return res; + } + + // TODO: we have no cases with multiple memory providers + assert(allocInfo.pool->numProviders == 1); + umf_memory_provider_handle_t provider = allocInfo.pool->providers[0]; + + size_t providerIPCHandleSize; + res = umfMemoryProviderGetIPCHandleSize(provider, &providerIPCHandleSize); + if (res != UMF_RESULT_SUCCESS) { + return res; + } + + ipcHandleSize = sizeof(struct umf_ipc_data_t) + providerIPCHandleSize; + struct umf_ipc_data_t *ipcData = malloc(ipcHandleSize); + if (!ipcData) { + return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY; + } + + res = + umfMemoryProviderGetIPCHandle(provider, allocInfo.base, allocInfo.size, + (void *)ipcData->providerData); + if (res != UMF_RESULT_SUCCESS) { + free(ipcData); + return res; + } + + ipcData->size = ipcHandleSize; + ipcData->offset = (uintptr_t)ptr - (uintptr_t)allocInfo.base; + + *umfIPCHandle = ipcData; + *size = ipcHandleSize; + + return UMF_RESULT_SUCCESS; +} + +enum umf_result_t umfPutIPCHandle(umf_ipc_handle_t umfIPCHandle) { + return UMF_RESULT_ERROR_NOT_SUPPORTED; +} + +enum umf_result_t umfOpenIPCHandle(umf_memory_pool_handle_t hPool, + umf_ipc_handle_t umfIPCHandle, void **ptr) { + umf_memory_provider_handle_t hProvider; + size_t numProviders; + void *base = NULL; + + // TODO: So far we have no pools that support more then 1 memory providers. + enum umf_result_t res = + umfPoolGetMemoryProviders(hPool, 1, &hProvider, &numProviders); + if (res != UMF_RESULT_SUCCESS) { + return res; + } + + umfMemoryProviderOpenIPCHandle(hProvider, + (void *)umfIPCHandle->providerData, &base); + + *ptr = (void *)((uintptr_t)base + umfIPCHandle->offset); + + return UMF_RESULT_SUCCESS; +} + +enum umf_result_t umfCloseIPCHandle(void *ptr) { + return UMF_RESULT_ERROR_NOT_SUPPORTED; + umf_memory_provider_handle_t hProvider = NULL; // TODO: find memory provider + + return umfMemoryProviderCloseIPCHandle(hProvider, ptr); +} diff --git a/source/common/unified_malloc_framework/src/memory_provider.cpp b/source/common/unified_malloc_framework/src/memory_provider.cpp index 084c311767..14f7b808d0 100644 --- a/source/common/unified_malloc_framework/src/memory_provider.cpp +++ b/source/common/unified_malloc_framework/src/memory_provider.cpp @@ -171,6 +171,32 @@ const char *umfMemoryProviderGetName(umf_memory_provider_handle_t hProvider) { return hProvider->ops.get_name(hProvider->provider_priv); } +enum umf_result_t +umfMemoryProviderGetIPCHandleSize(umf_memory_provider_handle_t hProvider, + size_t *size) { + return hProvider->ops.get_ipc_handle_size(hProvider->provider_priv, size); +} + +enum umf_result_t +umfMemoryProviderGetIPCHandle(umf_memory_provider_handle_t hProvider, + const void *ptr, size_t size, void *ipcData) { + return hProvider->ops.get_ipc_handle(hProvider->provider_priv, ptr, size, + ipcData); +} + +enum umf_result_t +umfMemoryProviderOpenIPCHandle(umf_memory_provider_handle_t hProvider, + void *ipcData, void **ptr) { + return hProvider->ops.open_ipc_handle(hProvider->provider_priv, ipcData, + ptr); +} + +enum umf_result_t +umfMemoryProviderCloseIPCHandle(umf_memory_provider_handle_t hProvider, + void *ptr) { + return hProvider->ops.close_ipc_handle(hProvider->provider_priv, ptr); +} + umf_memory_provider_handle_t umfGetLastFailedMemoryProvider(void) { return *umfGetLastFailedMemoryProviderPtr(); } diff --git a/source/common/unified_malloc_framework/src/memory_tracker.cpp b/source/common/unified_malloc_framework/src/memory_tracker.cpp index adbe2aa5e9..6c4baa13f5 100644 --- a/source/common/unified_malloc_framework/src/memory_tracker.cpp +++ b/source/common/unified_malloc_framework/src/memory_tracker.cpp @@ -22,6 +22,10 @@ #include #endif +struct tracker_value_t { + size_t size; + void *pool; +}; // TODO: reimplement in C and optimize... struct umf_memory_tracker_t { enum umf_result_t add(void *pool, const void *ptr, size_t size) { @@ -31,8 +35,8 @@ struct umf_memory_tracker_t { return UMF_RESULT_SUCCESS; } - auto ret = - map.try_emplace(reinterpret_cast(ptr), size, pool); + auto ret = map.try_emplace(reinterpret_cast(ptr), + tracker_value_t{size, pool}); return ret.second ? UMF_RESULT_SUCCESS : UMF_RESULT_ERROR_UNKNOWN; } @@ -47,31 +51,34 @@ struct umf_memory_tracker_t { return UMF_RESULT_SUCCESS; } - void *find(const void *ptr) { + bool find(const void *ptr, umf_alloc_info_t *pAllocInfo) { std::shared_lock lock(mtx); auto intptr = reinterpret_cast(ptr); auto it = map.upper_bound(intptr); if (it == map.begin()) { - return nullptr; + return false; } --it; auto address = it->first; - auto size = it->second.first; - auto pool = it->second.second; + auto size = it->second.size; + auto pool = it->second.pool; if (intptr >= address && intptr < address + size) { - return pool; + pAllocInfo->base = reinterpret_cast(address); + pAllocInfo->size = size; + pAllocInfo->pool = (umf_memory_pool_handle_t)pool; + return true; } - return nullptr; + return false; } private: std::shared_mutex mtx; - std::map> map; + std::map map; }; static enum umf_result_t @@ -114,7 +121,15 @@ umf_memory_tracker_handle_t umfMemoryTrackerGet(void) { return tracker; } void *umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, const void *ptr) { - return hTracker->find(ptr); + struct umf_alloc_info_t allocInfo; + return hTracker->find(ptr, &allocInfo) ? allocInfo.pool : nullptr; +} + +enum umf_result_t +umfMemoryTrackerGetAllocInfo(umf_memory_tracker_handle_t hTracker, + const void *ptr, umf_alloc_info_t *pAllocInfo) { + return hTracker->find(ptr, pAllocInfo) ? UMF_RESULT_SUCCESS + : UMF_RESULT_ERROR_INVALID_ARGUMENT; } struct umf_tracking_memory_provider_t { @@ -233,6 +248,33 @@ static const char *trackingName(void *provider) { return umfMemoryProviderGetName(p->hUpstream); } +static enum umf_result_t trackingGetIpcHandleSize(void *provider, + size_t *size) { + umf_tracking_memory_provider_t *p = + (umf_tracking_memory_provider_t *)provider; + return umfMemoryProviderGetIPCHandleSize(p->hUpstream, size); +} + +static enum umf_result_t trackingGetIpcHandle(void *provider, const void *ptr, + size_t size, void *ipcData) { + umf_tracking_memory_provider_t *p = + (umf_tracking_memory_provider_t *)provider; + return umfMemoryProviderGetIPCHandle(p->hUpstream, ptr, size, ipcData); +} + +static enum umf_result_t trackingOpenIpcHandle(void *provider, void *ipcData, + void **ptr) { + umf_tracking_memory_provider_t *p = + (umf_tracking_memory_provider_t *)provider; + return umfMemoryProviderOpenIPCHandle(p->hUpstream, ipcData, ptr); +} + +static enum umf_result_t trackingCloseIpcHandle(void *provider, void *ptr) { + umf_tracking_memory_provider_t *p = + (umf_tracking_memory_provider_t *)provider; + return umfMemoryProviderCloseIPCHandle(p->hUpstream, ptr); +} + enum umf_result_t umfTrackingMemoryProviderCreate( umf_memory_provider_handle_t hUpstream, umf_memory_pool_handle_t hPool, umf_memory_provider_handle_t *hTrackingProvider) { @@ -254,6 +296,10 @@ enum umf_result_t umfTrackingMemoryProviderCreate( trackingMemoryProviderOps.purge_force = trackingPurgeForce; trackingMemoryProviderOps.purge_lazy = trackingPurgeLazy; trackingMemoryProviderOps.get_name = trackingName; + trackingMemoryProviderOps.get_ipc_handle_size = trackingGetIpcHandleSize; + trackingMemoryProviderOps.get_ipc_handle = trackingGetIpcHandle; + trackingMemoryProviderOps.open_ipc_handle = trackingOpenIpcHandle; + trackingMemoryProviderOps.close_ipc_handle = trackingCloseIpcHandle; return umfMemoryProviderCreate(&trackingMemoryProviderOps, ¶ms, hTrackingProvider); diff --git a/source/common/unified_malloc_framework/src/memory_tracker.h b/source/common/unified_malloc_framework/src/memory_tracker.h index 43a95cf0cd..2e5cfb85c6 100644 --- a/source/common/unified_malloc_framework/src/memory_tracker.h +++ b/source/common/unified_malloc_framework/src/memory_tracker.h @@ -21,10 +21,21 @@ extern "C" { typedef struct umf_memory_tracker_t *umf_memory_tracker_handle_t; +struct umf_alloc_info_t { + void *base; + size_t size; + umf_memory_pool_handle_t pool; +}; + umf_memory_tracker_handle_t umfMemoryTrackerGet(void); void *umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, const void *ptr); +enum umf_result_t +umfMemoryTrackerGetAllocInfo(umf_memory_tracker_handle_t hTracker, + const void *ptr, + struct umf_alloc_info_t *pAllocInfo); + // Creates a memory provider that tracks each allocation/deallocation through umf_memory_tracker_handle_t and // forwards all requests to hUpstream memory Provider. hUpstream lifetime should be managed by the user of this function. enum umf_result_t umfTrackingMemoryProviderCreate(