Merge pull request #443 from igchor/get_providers

[uma] add function to retrieve providers

Author: pbalcer

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

@@ -129,6 +129,18 @@ enum uma_result_t umaPoolGetLastResult(uma_memory_pool_handle_t hPool,
 /// \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);
 
+///
+/// \brief Retrieve memory providers associated with a given pool.
+/// \param hPool specified memory pool
+/// \param hProviders [out] pointer to an array of memory providers. If numProviders is not equal to or
+///        greater than the real number of providers, UMA_RESULT_ERROR_INVALID_ARGUMENT is returned.
+/// \param numProviders [in] number of memory providers to return
+/// \param numProvidersRet pointer to the actual number of memory providers.
+enum uma_result_t
+umaPoolGetMemoryProviders(uma_memory_pool_handle_t hPool, size_t numProviders,
+                          uma_memory_provider_handle_t *hProviders,
+                          size_t *numProvidersRet);
+
 #ifdef __cplusplus
 }
 #endif

source/common/unified_memory_allocation/src/memory_pool.c

@@ -6,7 +6,9 @@
  *
  */
 
+#include "memory_provider_internal.h"
 #include "memory_tracker.h"
+
 #include <uma/memory_pool.h>
 #include <uma/memory_pool_ops.h>
 
@@ -134,3 +136,25 @@ enum uma_result_t umaPoolGetLastResult(uma_memory_pool_handle_t hPool,
 uma_memory_pool_handle_t umaPoolByPtr(const void *ptr) {
     return umaMemoryTrackerGetPool(umaMemoryTrackerGet(), ptr);
 }
+
+enum uma_result_t
+umaPoolGetMemoryProviders(uma_memory_pool_handle_t hPool, size_t numProviders,
+                          uma_memory_provider_handle_t *hProviders,
+                          size_t *numProvidersRet) {
+    if (hProviders && numProviders < hPool->numProviders) {
+        return UMA_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    if (numProvidersRet) {
+        *numProvidersRet = hPool->numProviders;
+    }
+
+    if (hProviders) {
+        for (size_t i = 0; i < hPool->numProviders; i++) {
+            umaTrackingMemoryProviderGetUpstreamProvider(
+                umaMemoryProviderGetPriv(hPool->providers[i]), hProviders + i);
+        }
+    }
+
+    return UMA_RESULT_SUCCESS;
+}

source/common/unified_memory_allocation/src/memory_provider.c

@@ -6,6 +6,7 @@
  *
  */
 
+#include "memory_provider_internal.h"
 #include <uma/memory_provider.h>
 
 #include <assert.h>
@@ -64,3 +65,7 @@ umaMemoryProviderGetLastResult(uma_memory_provider_handle_t hProvider,
                                const char **ppMessage) {
     return hProvider->ops.get_last_result(hProvider->provider_priv, ppMessage);
 }
+
+void *umaMemoryProviderGetPriv(uma_memory_provider_handle_t hProvider) {
+    return hProvider->provider_priv;
+}

source/common/unified_memory_allocation/src/memory_provider_internal.h

@@ -0,0 +1,24 @@
+/*
+ *
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#ifndef UMA_MEMORY_PROVIDER_INTERNAL_H
+#define UMA_MEMORY_PROVIDER_INTERNAL_H 1
+
+#include <uma/memory_provider.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+void *umaMemoryProviderGetPriv(uma_memory_provider_handle_t hProvider);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* UMA_MEMORY_PROVIDER_INTERNAL_H */

source/common/unified_memory_allocation/src/memory_tracker.cpp

@@ -10,6 +10,7 @@
 #include <uma/memory_provider.h>
 #include <uma/memory_provider_ops.h>
 
+#include <cassert>
 #include <map>
 #include <mutex>
 #include <shared_mutex>
@@ -171,7 +172,8 @@ enum uma_result_t umaTrackingMemoryProviderCreate(
     uma_memory_provider_handle_t hUpstream, uma_memory_pool_handle_t hPool,
     uma_memory_provider_handle_t *hTrackingProvider) {
     uma_tracking_memory_provider_t params;
-    params.hUpstream = hUpstream, params.hTracker = umaMemoryTrackerGet(),
+    params.hUpstream = hUpstream;
+    params.hTracker = umaMemoryTrackerGet();
     params.pool = hPool;
 
     struct uma_memory_provider_ops_t trackingMemoryProviderOps;
@@ -185,4 +187,13 @@ enum uma_result_t umaTrackingMemoryProviderCreate(
     return umaMemoryProviderCreate(&trackingMemoryProviderOps, &params,
                                    hTrackingProvider);
 }
+
+void umaTrackingMemoryProviderGetUpstreamProvider(
+    uma_memory_provider_handle_t hTrackingProvider,
+    uma_memory_provider_handle_t *hUpstream) {
+    assert(hUpstream);
+    uma_tracking_memory_provider_t *p =
+        (uma_tracking_memory_provider_t *)hTrackingProvider;
+    *hUpstream = p->hUpstream;
+}
 }

source/common/unified_memory_allocation/src/memory_tracker.h

@@ -29,6 +29,10 @@ enum uma_result_t umaTrackingMemoryProviderCreate(
     uma_memory_provider_handle_t hUpstream, uma_memory_pool_handle_t hPool,
     uma_memory_provider_handle_t *hTrackingProvider);
 
+void umaTrackingMemoryProviderGetUpstreamProvider(
+    uma_memory_provider_handle_t hTrackingProvider,
+    uma_memory_provider_handle_t *hUpstream);
+
 #ifdef __cplusplus
 }
 #endif

test/unified_memory_allocation/common/pool.hpp

@@ -69,7 +69,6 @@ struct malloc_pool : public pool_base {
 struct proxy_pool : public pool_base {
     uma_result_t initialize(uma_memory_provider_handle_t *providers,
                             size_t numProviders) noexcept {
-        EXPECT_EQ(numProviders, 1);
         this->provider = providers[0];
         return UMA_RESULT_SUCCESS;
     }

test/unified_memory_allocation/memoryPoolAPI.cpp

@@ -9,6 +9,7 @@
 
 #include "memoryPool.hpp"
 
+#include <array>
 #include <string>
 #include <unordered_map>
 
@@ -107,6 +108,28 @@ TEST_F(test, memoryPoolWithCustomProviders) {
     }
 }
 
+TEST_F(test, retrieveMemoryProviders) {
+    static constexpr size_t numProviders = 4;
+    std::array<uma_memory_provider_handle_t, numProviders> providers = {
+        (uma_memory_provider_handle_t)0x1, (uma_memory_provider_handle_t)0x2,
+        (uma_memory_provider_handle_t)0x3, (uma_memory_provider_handle_t)0x4};
+
+    auto [ret, pool] = uma::poolMakeUnique<uma_test::proxy_pool>(
+        providers.data(), numProviders);
+
+    std::array<uma_memory_provider_handle_t, numProviders> retProviders;
+    size_t numProvidersRet = 0;
+
+    ret = umaPoolGetMemoryProviders(pool.get(), 0, nullptr, &numProvidersRet);
+    ASSERT_EQ(ret, UMA_RESULT_SUCCESS);
+    ASSERT_EQ(numProvidersRet, numProviders);
+
+    ret = umaPoolGetMemoryProviders(pool.get(), numProviders,
+                                    retProviders.data(), nullptr);
+    ASSERT_EQ(ret, UMA_RESULT_SUCCESS);
+    ASSERT_EQ(retProviders, providers);
+}
+
 template <typename Pool>
 static auto
 makePool(std::function<uma::provider_unique_handle_t()> makeProvider) {
@@ -185,3 +208,16 @@ TEST_P(poolInitializeTest, errorPropagation) {
     ASSERT_EQ(ret.first, this->GetParam());
     ASSERT_EQ(ret.second, nullptr);
 }
+
+TEST_F(test, retrieveMemoryProvidersError) {
+    static constexpr size_t numProviders = 4;
+    std::array<uma_memory_provider_handle_t, numProviders> providers = {
+        (uma_memory_provider_handle_t)0x1, (uma_memory_provider_handle_t)0x2,
+        (uma_memory_provider_handle_t)0x3, (uma_memory_provider_handle_t)0x4};
+
+    auto [ret, pool] = uma::poolMakeUnique<uma_test::proxy_pool>(
+        providers.data(), numProviders);
+
+    ret = umaPoolGetMemoryProviders(pool.get(), 1, providers.data(), nullptr);
+    ASSERT_EQ(ret, UMA_RESULT_ERROR_INVALID_ARGUMENT);
+}