Merge pull request #315 from igchor/uma_memory_tracking_provider

[uma] memory tracking provider

Author: pbalcer

source/common/uma_helpers.hpp

@@ -14,6 +14,7 @@
 #include <uma/memory_provider.h>
 #include <uma/memory_provider_ops.h>
 
+#include <functional>
 #include <memory>
 #include <stdexcept>
 #include <tuple>
@@ -22,9 +23,11 @@
 namespace uma {
 
 using pool_unique_handle_t =
-    std::unique_ptr<uma_memory_pool_t, decltype(&umaPoolDestroy)>;
+    std::unique_ptr<uma_memory_pool_t,
+                    std::function<void(uma_memory_pool_handle_t)>>;
 using provider_unique_handle_t =
-    std::unique_ptr<uma_memory_provider_t, decltype(&umaMemoryProviderDestroy)>;
+    std::unique_ptr<uma_memory_provider_t,
+                    std::function<void(uma_memory_provider_handle_t)>>;
 
 /// @brief creates UMA memory provider based on given T type.
 /// T should implement all functions defined by
@@ -45,7 +48,7 @@ auto memoryProviderMakeUnique(Args &&...args) {
         auto provider = new T;
         *obj = provider;
         return std::apply(&T::initialize,
-                          std::tuple_cat(std::make_tuple(*provider), *tuple));
+                          std::tuple_cat(std::make_tuple(provider), *tuple));
     };
     ops.finalize = [](void *obj) { delete reinterpret_cast<T *>(obj); };
     ops.alloc = [](void *obj, auto... args) {
@@ -74,25 +77,34 @@ auto memoryProviderMakeUnique(Args &&...args) {
 /// replaced by dtor). All arguments passed to this function are
 /// forwarded to T::initialize(). All functions of T
 /// should be noexcept.
-template <typename T, typename... Args> auto poolMakeUnique(Args &&...args) {
+template <typename T, typename... Args>
+auto poolMakeUnique(uma_memory_provider_handle_t *providers,
+                    size_t numProviders, Args &&...args) {
     uma_memory_pool_ops_t ops;
     auto argsTuple = std::make_tuple(std::forward<Args>(args)...);
-    static_assert(
-        noexcept(std::declval<T>().initialize(std::forward<Args>(args)...)));
+    static_assert(noexcept(std::declval<T>().initialize(
+        providers, numProviders, std::forward<Args>(args)...)));
 
     ops.version = UMA_VERSION_CURRENT;
-    ops.initialize = [](void *params, void **obj) {
+    ops.initialize = [](uma_memory_provider_handle_t *providers,
+                        size_t numProviders, void *params, void **obj) {
         auto *tuple = reinterpret_cast<decltype(argsTuple) *>(params);
         auto pool = new T;
         *obj = pool;
-        return std::apply(&T::initialize,
-                          std::tuple_cat(std::make_tuple(*pool), *tuple));
+        return std::apply(
+            &T::initialize,
+            std::tuple_cat(std::make_tuple(pool, providers, numProviders),
+                           *tuple));
     };
     ops.finalize = [](void *obj) { delete reinterpret_cast<T *>(obj); };
     ops.malloc = [](void *obj, auto... args) {
         static_assert(noexcept(reinterpret_cast<T *>(obj)->malloc(args...)));
         return reinterpret_cast<T *>(obj)->malloc(args...);
     };
+    ops.calloc = [](void *obj, auto... args) {
+        static_assert(noexcept(reinterpret_cast<T *>(obj)->calloc(args...)));
+        return reinterpret_cast<T *>(obj)->calloc(args...);
+    };
     ops.aligned_malloc = [](void *obj, auto... args) {
         static_assert(
             noexcept(reinterpret_cast<T *>(obj)->aligned_malloc(args...)));
@@ -118,7 +130,7 @@ template <typename T, typename... Args> auto poolMakeUnique(Args &&...args) {
     };
 
     uma_memory_pool_handle_t hPool = nullptr;
-    auto ret = umaPoolCreate(&ops, &argsTuple, &hPool);
+    auto ret = umaPoolCreate(&ops, providers, numProviders, &argsTuple, &hPool);
     return std::pair<uma_result_t, pool_unique_handle_t>{
         ret, pool_unique_handle_t(hPool, &umaPoolDestroy)};
 }

source/common/unified_memory_allocation/CMakeLists.txt

@@ -4,6 +4,7 @@
 set(UMA_SOURCES
     src/memory_pool.c
     src/memory_provider.c
+    src/memory_tracker.cpp
 )
 
 if(UMA_BUILD_SHARED_LIBRARY)

source/common/unified_memory_allocation/include/uma/memory_pool.h

@@ -10,21 +10,27 @@
 #define UMA_MEMORY_POOL_H 1
 
 #include <uma/base.h>
-#include <uma/memory_pool_ops.h>
+#include <uma/memory_provider.h>
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 typedef struct uma_memory_pool_t *uma_memory_pool_handle_t;
 
+struct uma_memory_pool_ops_t;
+
 ///
 /// \brief Creates new memory pool
 /// \param ops instance of uma_memory_pool_ops_t
+/// \param providers array of memory providers that will be used for coarse-grain allocations. Should contain at least one memory provider.
+/// \param numProvider number of elements in the providers array
 /// \param params pointer to pool-specific parameters
 /// \return UMA_RESULT_SUCCESS on success or appropriate error code on failure
 ///
-enum uma_result_t umaPoolCreate(struct uma_memory_pool_ops_t *ops, void *params,
+enum uma_result_t umaPoolCreate(struct uma_memory_pool_ops_t *ops,
+                                uma_memory_provider_handle_t *providers,
+                                size_t numProviders, void *params,
                                 uma_memory_pool_handle_t *hPool);
 
 ///
@@ -87,6 +93,12 @@ size_t umaPoolMallocUsableSize(uma_memory_pool_handle_t hPool, void *ptr);
 ///
 void umaPoolFree(uma_memory_pool_handle_t hPool, void *ptr);
 
+///
+/// \brief Frees the memory space pointed by ptr if it belongs to UMA pool, does nothing otherwise
+/// \param ptr pointer to the allocated memory
+///
+void umaFree(void *ptr);
+
 ///
 /// \brief Retrieve string representation of the underlying pool specific
 ///        result reported by the last API that returned
@@ -111,6 +123,12 @@ void umaPoolFree(uma_memory_pool_handle_t hPool, void *ptr);
 enum uma_result_t umaPoolGetLastResult(uma_memory_pool_handle_t hPool,
                                        const char **ppMessage);
 
+///
+/// \brief Retrieve memory pool associated with a given ptr.
+/// \param ptr pointer to memory belonging to a memory pool
+/// \return handle to a memory pool that contains ptr or NULL if pointer does not belong to any UMA pool
+uma_memory_pool_handle_t umaPoolByPtr(const void *ptr);
+
 #ifdef __cplusplus
 }
 #endif

source/common/unified_memory_allocation/include/uma/memory_pool_ops.h

@@ -15,7 +15,7 @@
 extern "C" {
 #endif
 
-/// This structure comprises function pointers used by corresponding  umaPool*
+/// \brief This structure comprises function pointers used by corresponding  umaPool*
 /// calls. Each memory pool implementation should initialize all function
 /// pointers.
 struct uma_memory_pool_ops_t {
@@ -25,11 +25,15 @@ struct uma_memory_pool_ops_t {
 
     ///
     /// \brief Intializes memory pool
+    /// \param providers array of memory providers that will be used for coarse-grain allocations. Should contain at least one memory provider.
+    /// \param numProvider number of elements in the providers array
     /// \param params pool-specific params
     /// \param pool returns pointer to the pool
     /// \return UMA_RESULT_SUCCESS on success or appropriate error code on
     /// failure
-    enum uma_result_t (*initialize)(void *params, void **pool);
+    enum uma_result_t (*initialize)(uma_memory_provider_handle_t *providers,
+                                    size_t numProviders, void *params,
+                                    void **pool);
 
     ///
     /// \brief Finalizes memory pool

source/common/unified_memory_allocation/include/uma/memory_provider.h

@@ -41,7 +41,7 @@ void umaMemoryProviderDestroy(uma_memory_provider_handle_t hProvider);
 /// \param hProvider handle to the memory provider
 /// \param size number of bytes to allocate
 /// \param alignment alignment of the allocation
-/// \param ptr returns pointer to the allocated memory
+/// \param ptr will be updated with pointer to the allocated memory
 /// \return UMA_RESULT_SUCCESS on success or appropriate error code on failure
 ///
 enum uma_result_t umaMemoryProviderAlloc(uma_memory_provider_handle_t hProvider,

source/common/unified_memory_allocation/src/memory_pool.c

@@ -6,43 +6,91 @@
  *
  */
 
+#include "memory_tracker.h"
 #include <uma/memory_pool.h>
+#include <uma/memory_pool_ops.h>
 
 #include <assert.h>
 #include <stdlib.h>
 
 struct uma_memory_pool_t {
-    struct uma_memory_pool_ops_t ops;
     void *pool_priv;
+    struct uma_memory_pool_ops_t ops;
+
+    // Holds array of memory providers. All providers are wrapped
+    // by memory tracking providers (owned and released by UMA).
+    uma_memory_provider_handle_t *providers;
+
+    size_t numProviders;
 };
 
-enum uma_result_t umaPoolCreate(struct uma_memory_pool_ops_t *ops, void *params,
+static void
+destroyMemoryProviderWrappers(uma_memory_provider_handle_t *providers,
+                              size_t numProviders) {
+    for (size_t i = 0; i < numProviders; i++) {
+        umaMemoryProviderDestroy(providers[i]);
+    }
+
+    free(providers);
+}
+
+enum uma_result_t umaPoolCreate(struct uma_memory_pool_ops_t *ops,
+                                uma_memory_provider_handle_t *providers,
+                                size_t numProviders, void *params,
                                 uma_memory_pool_handle_t *hPool) {
+    if (!numProviders || !providers) {
+        return UMA_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    enum uma_result_t ret = UMA_RESULT_SUCCESS;
     uma_memory_pool_handle_t pool = malloc(sizeof(struct uma_memory_pool_t));
     if (!pool) {
         return UMA_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
     assert(ops->version == UMA_VERSION_CURRENT);
 
-    pool->ops = *ops;
+    pool->providers =
+        calloc(numProviders, sizeof(uma_memory_provider_handle_t));
+    if (!pool->providers) {
+        ret = UMA_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+        goto err_providers_alloc;
+    }
 
-    void *pool_priv;
-    enum uma_result_t ret = ops->initialize(params, &pool_priv);
-    if (ret != UMA_RESULT_SUCCESS) {
-        free(pool);
-        return ret;
+    size_t providerInd = 0;
+    pool->numProviders = numProviders;
+
+    // Wrap each provider with memory tracking provider.
+    for (providerInd = 0; providerInd < numProviders; providerInd++) {
+        ret = umaTrackingMemoryProviderCreate(providers[providerInd], pool,
+                                              &pool->providers[providerInd]);
+        if (ret != UMA_RESULT_SUCCESS) {
+            goto err_providers_init;
+        }
     }
 
-    pool->pool_priv = pool_priv;
+    pool->ops = *ops;
+    ret = ops->initialize(pool->providers, pool->numProviders, params,
+                          &pool->pool_priv);
+    if (ret != UMA_RESULT_SUCCESS) {
+        goto err_pool_init;
+    }
 
     *hPool = pool;
-
     return UMA_RESULT_SUCCESS;
+
+err_pool_init:
+err_providers_init:
+    destroyMemoryProviderWrappers(pool->providers, providerInd);
+err_providers_alloc:
+    free(pool);
+
+    return ret;
 }
 
 void umaPoolDestroy(uma_memory_pool_handle_t hPool) {
     hPool->ops.finalize(hPool->pool_priv);
+    destroyMemoryProviderWrappers(hPool->providers, hPool->numProviders);
     free(hPool);
 }
 
@@ -71,7 +119,18 @@ void umaPoolFree(uma_memory_pool_handle_t hPool, void *ptr) {
     hPool->ops.free(hPool->pool_priv, ptr);
 }
 
+void umaFree(void *ptr) {
+    uma_memory_pool_handle_t hPool = umaPoolByPtr(ptr);
+    if (hPool) {
+        umaPoolFree(hPool, ptr);
+    }
+}
+
 enum uma_result_t umaPoolGetLastResult(uma_memory_pool_handle_t hPool,
                                        const char **ppMessage) {
     return hPool->ops.get_last_result(hPool->pool_priv, ppMessage);
 }
+
+uma_memory_pool_handle_t umaPoolByPtr(const void *ptr) {
+    return umaMemoryTrackerGetPool(umaMemoryTrackerGet(), ptr);
+}