Skip to content

Commit

Permalink
Add IPC API to support IPC handles
Browse files Browse the repository at this point in the history
  • Loading branch information
vinser52 committed Sep 19, 2023
1 parent 92d57c3 commit f96e337
Show file tree
Hide file tree
Showing 8 changed files with 294 additions and 10 deletions.
1 change: 1 addition & 0 deletions source/common/unified_malloc_framework/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
56 changes: 56 additions & 0 deletions source/common/unified_malloc_framework/include/umf/ipc.h
Original file line number Diff line number Diff line change
@@ -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 <umf/base.h>
#include <umf/memory_pool.h>

#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 */
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
98 changes: 98 additions & 0 deletions source/common/unified_malloc_framework/src/ipc.c
Original file line number Diff line number Diff line change
@@ -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 <assert.h>
#include <stdlib.h>

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);
}
26 changes: 26 additions & 0 deletions source/common/unified_malloc_framework/src/memory_provider.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
66 changes: 56 additions & 10 deletions source/common/unified_malloc_framework/src/memory_tracker.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@
#include <windows.h>
#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) {
Expand All @@ -31,8 +35,8 @@ struct umf_memory_tracker_t {
return UMF_RESULT_SUCCESS;
}

auto ret =
map.try_emplace(reinterpret_cast<uintptr_t>(ptr), size, pool);
auto ret = map.try_emplace(reinterpret_cast<uintptr_t>(ptr),
tracker_value_t{size, pool});
return ret.second ? UMF_RESULT_SUCCESS : UMF_RESULT_ERROR_UNKNOWN;
}

Expand All @@ -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<std::shared_mutex> lock(mtx);

auto intptr = reinterpret_cast<uintptr_t>(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<void *>(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<uintptr_t, std::pair<size_t, void *>> map;
std::map<uintptr_t, tracker_value_t> map;
};

static enum umf_result_t
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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) {
Expand All @@ -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, &params,
hTrackingProvider);
Expand Down
Loading

0 comments on commit f96e337

Please sign in to comment.