diff --git a/source/common/unified_malloc_framework/CMakeLists.txt b/source/common/unified_malloc_framework/CMakeLists.txt index 15744605ec..384a83480b 100644 --- a/source/common/unified_malloc_framework/CMakeLists.txt +++ b/source/common/unified_malloc_framework/CMakeLists.txt @@ -6,10 +6,15 @@ set(UMF_SOURCES src/memory_pool.c src/memory_provider.c - src/memory_tracker.cpp + src/memory_tracker.c src/memory_provider_get_last_failed.cpp + src/critnib/critnib.c ) +if (MSVC) + set(UMF_SOURCES ${UMF_SOURCES} src/memory_tracker_windows.cpp) +endif() + if(UMF_BUILD_SHARED_LIBRARY) message(WARNING "Unified Malloc Framework is still an early work in progress." "There are no API/ABI backward compatibility guarantees. There will be breakages." diff --git a/source/common/unified_malloc_framework/src/memory_tracker.cpp b/source/common/unified_malloc_framework/src/memory_tracker.c similarity index 70% rename from source/common/unified_malloc_framework/src/memory_tracker.cpp rename to source/common/unified_malloc_framework/src/memory_tracker.c index adbe2aa5e9..55758adee9 100644 --- a/source/common/unified_malloc_framework/src/memory_tracker.cpp +++ b/source/common/unified_malloc_framework/src/memory_tracker.c @@ -9,112 +9,97 @@ */ #include "memory_tracker.h" +#include "critnib/critnib.h" + +#include #include #include -#include -#include -#include -#include +#include +#include #include -#ifdef _WIN32 -#include -#endif - -// TODO: reimplement in C and optimize... -struct umf_memory_tracker_t { - enum umf_result_t add(void *pool, const void *ptr, size_t size) { - std::unique_lock lock(mtx); +#if !defined(_WIN32) +critnib *TRACKER = NULL; +void __attribute__((constructor)) createLibTracker() { + TRACKER = critnib_new(); +} +void __attribute__((destructor)) deleteLibTracker() { critnib_delete(TRACKER); } - if (size == 0) { - return UMF_RESULT_SUCCESS; - } +umf_memory_tracker_handle_t umfMemoryTrackerGet(void) { + return (umf_memory_tracker_handle_t)TRACKER; +} +#endif - auto ret = - map.try_emplace(reinterpret_cast(ptr), size, pool); - return ret.second ? UMF_RESULT_SUCCESS : UMF_RESULT_ERROR_UNKNOWN; - } +struct tracker_value_t { + umf_memory_pool_handle_t pool; + size_t size; +}; - enum umf_result_t remove(const void *ptr, size_t size) { - std::unique_lock lock(mtx); +static enum umf_result_t +umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker, + umf_memory_pool_handle_t pool, const void *ptr, + size_t size) { + assert(ptr); - map.erase(reinterpret_cast(ptr)); + struct tracker_value_t *value = + (struct tracker_value_t *)malloc(sizeof(struct tracker_value_t)); + value->pool = pool; + value->size = size; - // TODO: handle removing part of the range - (void)size; + int ret = critnib_insert((critnib *)hTracker, (uintptr_t)ptr, value, 0); + if (ret == 0) { return UMF_RESULT_SUCCESS; } - void *find(const void *ptr) { - std::shared_lock lock(mtx); - - auto intptr = reinterpret_cast(ptr); - auto it = map.upper_bound(intptr); - if (it == map.begin()) { - return nullptr; - } - - --it; - - auto address = it->first; - auto size = it->second.first; - auto pool = it->second.second; - - if (intptr >= address && intptr < address + size) { - return pool; - } + free(value); - return nullptr; + if (ret == ENOMEM) { + return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY; } - private: - std::shared_mutex mtx; - std::map> map; -}; - -static enum umf_result_t -umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker, void *pool, - const void *ptr, size_t size) { - return hTracker->add(pool, ptr, size); + // This should not happen + // TODO: add logging here + return UMF_RESULT_ERROR_UNKNOWN; } static enum umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker, const void *ptr, size_t size) { - return hTracker->remove(ptr, size); -} + assert(ptr); + + // TODO: there is no support for removing partial ranges (or multipe entires + // in a single remove call) yet. + // Every umfMemoryTrackerAdd(..., ptr, ...) should have a corresponsding + // umfMemoryTrackerRemove call with the same ptr value. + (void)size; + + void *value = critnib_remove((critnib *)hTracker, (uintptr_t)ptr); + if (!value) { + // This should not happen + // TODO: add logging here + return UMF_RESULT_ERROR_UNKNOWN; + } -extern "C" { + free(value); -#if defined(_WIN32) && defined(UMF_SHARED_LIBRARY) -umf_memory_tracker_t *tracker = nullptr; -BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) { - if (fdwReason == DLL_PROCESS_DETACH) { - delete tracker; - } else if (fdwReason == DLL_PROCESS_ATTACH) { - tracker = new umf_memory_tracker_t; - } - return TRUE; -} -#elif defined(_WIN32) -umf_memory_tracker_t trackerInstance; -umf_memory_tracker_t *tracker = &trackerInstance; -#else -umf_memory_tracker_t *tracker = nullptr; -void __attribute__((constructor)) createLibTracker() { - tracker = new umf_memory_tracker_t; + return UMF_RESULT_SUCCESS; } -void __attribute__((destructor)) deleteLibTracker() { delete tracker; } -#endif +umf_memory_pool_handle_t +umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, const void *ptr) { + assert(ptr); -umf_memory_tracker_handle_t umfMemoryTrackerGet(void) { return tracker; } + uintptr_t rkey; + struct tracker_value_t *rvalue; + int found = critnib_find((critnib *)hTracker, (uintptr_t)ptr, FIND_LE, + (void *)&rkey, (void **)&rvalue); + if (!found) { + return NULL; + } -void *umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, - const void *ptr) { - return hTracker->find(ptr); + return (rkey + rvalue->size >= (uintptr_t)ptr) ? rvalue->pool : NULL; } struct umf_tracking_memory_provider_t { @@ -136,7 +121,7 @@ static enum umf_result_t trackingAlloc(void *hProvider, size_t size, } ret = umfMemoryProviderAlloc(p->hUpstream, size, alignment, ptr); - if (ret != UMF_RESULT_SUCCESS) { + if (ret != UMF_RESULT_SUCCESS || !*ptr) { return ret; } @@ -159,9 +144,11 @@ static enum umf_result_t trackingFree(void *hProvider, void *ptr, size_t size) { // to avoid a race condition. If the order would be different, other thread // could allocate the memory at address `ptr` before a call to umfMemoryTrackerRemove // resulting in inconsistent state. - ret = umfMemoryTrackerRemove(p->hTracker, ptr, size); - if (ret != UMF_RESULT_SUCCESS) { - return ret; + if (ptr) { + ret = umfMemoryTrackerRemove(p->hTracker, ptr, size); + if (ret != UMF_RESULT_SUCCESS) { + return ret; + } } ret = umfMemoryProviderFree(p->hUpstream, ptr, size); @@ -267,4 +254,3 @@ void umfTrackingMemoryProviderGetUpstreamProvider( (umf_tracking_memory_provider_t *)hTrackingProvider; *hUpstream = p->hUpstream; } -} diff --git a/source/common/unified_malloc_framework/src/memory_tracker.h b/source/common/unified_malloc_framework/src/memory_tracker.h index 43a95cf0cd..c16844928e 100644 --- a/source/common/unified_malloc_framework/src/memory_tracker.h +++ b/source/common/unified_malloc_framework/src/memory_tracker.h @@ -22,8 +22,8 @@ extern "C" { typedef struct umf_memory_tracker_t *umf_memory_tracker_handle_t; umf_memory_tracker_handle_t umfMemoryTrackerGet(void); -void *umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, - const void *ptr); +umf_memory_pool_handle_t +umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, const void *ptr); // 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. diff --git a/source/common/unified_malloc_framework/src/memory_tracker_windows.cpp b/source/common/unified_malloc_framework/src/memory_tracker_windows.cpp new file mode 100644 index 0000000000..036883301e --- /dev/null +++ b/source/common/unified_malloc_framework/src/memory_tracker_windows.cpp @@ -0,0 +1,35 @@ +/* + * + * 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 + +#if defined(UMF_SHARED_LIBRARY) +critnib *TRACKER = NULL; +BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) { + if (fdwReason == DLL_PROCESS_DETACH) { + critnib_delete(TRACKER); + } else if (fdwReason == DLL_PROCESS_ATTACH) { + TRACKER = critnib_new(); + } + return TRUE; +} +#else +struct tracker_t { + tracker_t() { map = critnib_new(); } + ~tracker_t() { critnib_remove(map); } + critnib *map; +}; +tracker_t TRACKER_INSTANCE; +critnib *TRACKER = TRACKER_INSTANCE.map; +#endif + +umf_memory_tracker_handle_t umfMemoryTrackerGet(void) { + return (umf_memory_tracker_handle_t)TRACKER; +} diff --git a/test/unified_malloc_framework/common/provider.c b/test/unified_malloc_framework/common/provider.c index 8f9e946bfc..303d8aea8d 100644 --- a/test/unified_malloc_framework/common/provider.c +++ b/test/unified_malloc_framework/common/provider.c @@ -23,7 +23,7 @@ static enum umf_result_t nullAlloc(void *provider, size_t size, (void)provider; (void)size; (void)alignment; - (void)ptr; + *ptr = NULL; return UMF_RESULT_SUCCESS; } diff --git a/test/unified_malloc_framework/memoryProviderAPI.cpp b/test/unified_malloc_framework/memoryProviderAPI.cpp index fa02f9eb99..02a7fa357d 100644 --- a/test/unified_malloc_framework/memoryProviderAPI.cpp +++ b/test/unified_malloc_framework/memoryProviderAPI.cpp @@ -23,7 +23,8 @@ TEST_F(test, memoryProviderTrace) { size_t call_count = 0; - auto ret = umfMemoryProviderAlloc(tracingProvider.get(), 0, 0, nullptr); + void *ptr; + auto ret = umfMemoryProviderAlloc(tracingProvider.get(), 0, 0, &ptr); ASSERT_EQ(ret, UMF_RESULT_SUCCESS); ASSERT_EQ(calls["alloc"], 1); ASSERT_EQ(calls.size(), ++call_count);