Add a basic pool manager for memory pools

Author: kswiecicki

source/common/ur_pool_manager.hpp

@@ -11,11 +11,17 @@
 #ifndef USM_POOL_MANAGER_HPP
 #define USM_POOL_MANAGER_HPP 1
 
+#include "logger/ur_logger.hpp"
+#include "umf_helpers.hpp"
+#include "umf_pools/disjoint_pool.hpp"
 #include "ur_api.h"
-#include "ur_pool_manager.hpp"
 #include "ur_util.hpp"
 
+#include <umf/memory_pool.h>
+#include <umf/memory_provider.h>
+
 #include <functional>
+#include <unordered_map>
 #include <vector>
 
 namespace usm {
@@ -29,8 +35,9 @@ struct pool_descriptor {
     ur_usm_type_t type;
     bool deviceReadOnly;
 
-    static bool equal(const pool_descriptor &lhs, const pool_descriptor &rhs);
-    static std::size_t hash(const pool_descriptor &desc);
+    bool operator==(const pool_descriptor &other) const;
+    friend std::ostream &operator<<(std::ostream &os,
+                                    const pool_descriptor &desc);
     static std::pair<ur_result_t, std::vector<pool_descriptor>>
     create(ur_usm_pool_handle_t poolHandle, ur_context_handle_t hContext);
 };
@@ -75,10 +82,10 @@ urGetSubDevices(ur_device_handle_t hDevice) {
 
 inline std::pair<ur_result_t, std::vector<ur_device_handle_t>>
 urGetAllDevicesAndSubDevices(ur_context_handle_t hContext) {
-    size_t deviceCount;
+    size_t deviceCount = 0;
     auto ret = urContextGetInfo(hContext, UR_CONTEXT_INFO_NUM_DEVICES,
                                 sizeof(deviceCount), &deviceCount, nullptr);
-    if (ret != UR_RESULT_SUCCESS) {
+    if (ret != UR_RESULT_SUCCESS || deviceCount == 0) {
         return {ret, {}};
     }
 
@@ -122,22 +129,28 @@ isSharedAllocationReadOnlyOnDevice(const pool_descriptor &desc) {
     return desc.type == UR_USM_TYPE_SHARED && desc.deviceReadOnly;
 }
 
-inline bool pool_descriptor::equal(const pool_descriptor &lhs,
-                                   const pool_descriptor &rhs) {
-    ur_native_handle_t lhsNative, rhsNative;
+inline bool pool_descriptor::operator==(const pool_descriptor &other) const {
+    const pool_descriptor &lhs = *this;
+    const pool_descriptor &rhs = other;
+    ur_native_handle_t lhsNative = nullptr, rhsNative = nullptr;
 
     // We want to share a memory pool for sub-devices and sub-sub devices.
     // Sub-devices and sub-sub-devices might be represented by different ur_device_handle_t but
     // they share the same native_handle_t (which is used by UMF provider).
     // Ref: https://github.com/intel/llvm/commit/86511c5dc84b5781dcfd828caadcb5cac157eae1
     // TODO: is this L0 specific?
-    auto ret = urDeviceGetNativeHandle(lhs.hDevice, &lhsNative);
-    if (ret != UR_RESULT_SUCCESS) {
-        throw ret;
+    if (lhs.hDevice) {
+        auto ret = urDeviceGetNativeHandle(lhs.hDevice, &lhsNative);
+        if (ret != UR_RESULT_SUCCESS) {
+            throw ret;
+        }
     }
-    ret = urDeviceGetNativeHandle(rhs.hDevice, &rhsNative);
-    if (ret != UR_RESULT_SUCCESS) {
-        throw ret;
+
+    if (rhs.hDevice) {
+        auto ret = urDeviceGetNativeHandle(rhs.hDevice, &rhsNative);
+        if (ret != UR_RESULT_SUCCESS) {
+            throw ret;
+        }
     }
 
     return lhsNative == rhsNative && lhs.type == rhs.type &&
@@ -146,16 +159,12 @@ inline bool pool_descriptor::equal(const pool_descriptor &lhs,
            lhs.poolHandle == rhs.poolHandle;
 }
 
-inline std::size_t pool_descriptor::hash(const pool_descriptor &desc) {
-    ur_native_handle_t native;
-    auto ret = urDeviceGetNativeHandle(desc.hDevice, &native);
-    if (ret != UR_RESULT_SUCCESS) {
-        throw ret;
-    }
-
-    return combine_hashes(0, desc.type, native,
-                          isSharedAllocationReadOnlyOnDevice(desc),
-                          desc.poolHandle);
+inline std::ostream &operator<<(std::ostream &os, const pool_descriptor &desc) {
+    os << "pool handle: " << desc.poolHandle
+       << " context handle: " << desc.hContext
+       << " device handle: " << desc.hDevice << " memory type: " << desc.type
+       << " is read only: " << desc.deviceReadOnly;
+    return os;
 }
 
 inline std::pair<ur_result_t, std::vector<pool_descriptor>>
@@ -177,6 +186,7 @@ pool_descriptor::create(ur_usm_pool_handle_t poolHandle,
             pool_descriptor &desc = descriptors.emplace_back();
             desc.poolHandle = poolHandle;
             desc.hContext = hContext;
+            desc.hDevice = device;
             desc.type = UR_USM_TYPE_DEVICE;
         }
         {
@@ -200,6 +210,69 @@ pool_descriptor::create(ur_usm_pool_handle_t poolHandle,
     return {ret, descriptors};
 }
 
+template <typename D> struct pool_manager {
+  private:
+    using desc_to_pool_map_t = std::unordered_map<D, umf::pool_unique_handle_t>;
+
+    desc_to_pool_map_t descToPoolMap;
+
+  public:
+    static std::pair<ur_result_t, pool_manager>
+    create(desc_to_pool_map_t descToHandleMap = {}) {
+        auto manager = pool_manager();
+
+        for (auto &[desc, hPool] : descToHandleMap) {
+            auto ret = manager.addPool(desc, hPool);
+            if (ret != UR_RESULT_SUCCESS) {
+                return {ret, pool_manager()};
+            }
+        }
+
+        return {UR_RESULT_SUCCESS, std::move(manager)};
+    }
+
+    ur_result_t addPool(const D &desc,
+                        umf::pool_unique_handle_t &hPool) noexcept {
+        if (!descToPoolMap.try_emplace(desc, std::move(hPool)).second) {
+            logger::error("Pool for pool descriptor: {}, already exists", desc);
+            return UR_RESULT_ERROR_INVALID_ARGUMENT;
+        }
+
+        return UR_RESULT_SUCCESS;
+    }
+
+    std::optional<umf_memory_pool_handle_t> getPool(const D &desc) noexcept {
+        auto it = descToPoolMap.find(desc);
+        if (it == descToPoolMap.end()) {
+            logger::error("Pool descriptor doesn't match any existing pool: {}",
+                          desc);
+            return std::nullopt;
+        }
+
+        return it->second.get();
+    }
+};
+
 } // namespace usm
 
+namespace std {
+/// @brief hash specialization for usm::pool_descriptor
+template <> struct hash<usm::pool_descriptor> {
+    inline size_t operator()(const usm::pool_descriptor &desc) const {
+        ur_native_handle_t native = nullptr;
+        if (desc.hDevice) {
+            auto ret = urDeviceGetNativeHandle(desc.hDevice, &native);
+            if (ret != UR_RESULT_SUCCESS) {
+                throw ret;
+            }
+        }
+
+        return combine_hashes(0, desc.type, native,
+                              isSharedAllocationReadOnlyOnDevice(desc),
+                              desc.poolHandle);
+    }
+};
+
+} // namespace std
+
 #endif /* USM_POOL_MANAGER_HPP */