Add a basic pool manager for memory pools

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);
 };
@@ -45,8 +52,8 @@ urGetSubDevices(ur_device_handle_t hDevice) {
  }
 
  ur_device_partition_property_t prop;
- prop.type = UR_DEVICE_PARTITION_EQUALLY;
- prop.value.equally = nComputeUnits;
+ prop.type = UR_DEVICE_PARTITION_BY_CSLICE;
+ prop.value.affinity_domain = 0;
 
  ur_device_partition_properties_t properties{
  UR_STRUCTURE_TYPE_DEVICE_PARTITION_PROPERTIES,
@@ -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, {}};
  }
 
@@ -110,6 +117,11 @@ urGetAllDevicesAndSubDevices(ur_context_handle_t hContext) {
  for (size_t i = 0; i < deviceCount; i++) {
  ret = addPoolsForDevicesRec(devices[i]);
  if (ret != UR_RESULT_SUCCESS) {
+ if (ret == UR_RESULT_ERROR_UNSUPPORTED_FEATURE) {
+ // Return main devices when sub-devices are unsupported.
+ return {ret, std::move(devices)};
+ }
+
  return {ret, {}};
  }
  }
@@ -122,22 +134,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 +164,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 +191,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 +215,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 */

test/usm/CMakeLists.txt

@@ -10,17 +10,21 @@ function(add_usm_test name)
  add_ur_executable(${TEST_TARGET_NAME}
  ${UR_USM_TEST_DIR}/../conformance/source/environment.cpp
  ${UR_USM_TEST_DIR}/../conformance/source/main.cpp
+ ${UR_USM_TEST_DIR}/../unified_malloc_framework/common/provider.c
+ ${UR_USM_TEST_DIR}/../unified_malloc_framework/common/pool.c
  ${ARGN})
  target_link_libraries(${TEST_TARGET_NAME}
  PRIVATE
  ${PROJECT_NAME}::common
  ${PROJECT_NAME}::loader
  ur_testing
  GTest::gtest_main)
- add_test(NAME  usm-${name}
+ add_test(NAME usm-${name}
  COMMAND ${TEST_TARGET_NAME}
  WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
- set_tests_properties(usm-${name} PROPERTIES LABELS "usm")
+ set_tests_properties(usm-${name} PROPERTIES
+ LABELS "usm"
+ ENVIRONMENT "UR_ADAPTERS_FORCE_LOAD=\"$<TARGET_FILE:ur_adapter_null>\"")
  target_compile_definitions("usm_test-${name}" PRIVATE DEVICES_ENVIRONMENT)
 endfunction()
 

test/usm/usmPoolManager.cpp

@@ -3,19 +3,18 @@
 // See LICENSE.TXT
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "../unified_malloc_framework/common/pool.hpp"
-#include "../unified_malloc_framework/common/provider.hpp"
 #include "ur_pool_manager.hpp"
 
-#include <uur/fixtures.h>
+#include "../unified_malloc_framework/common/pool.h"
+#include "../unified_malloc_framework/common/provider.h"
 
-#include <unordered_set>
+#include <uur/fixtures.h>
 
-struct urUsmPoolManagerTest
+struct urUsmPoolDescriptorTest
  : public uur::urMultiDeviceContextTest,
  ::testing::WithParamInterface<ur_usm_pool_handle_t> {};
 
-TEST_P(urUsmPoolManagerTest, poolIsPerContextTypeAndDevice) {
+TEST_P(urUsmPoolDescriptorTest, poolIsPerContextTypeAndDevice) {
  auto &devices = uur::DevicesEnvironment::instance->devices;
  auto poolHandle = this->GetParam();
 
@@ -49,7 +48,71 @@ TEST_P(urUsmPoolManagerTest, poolIsPerContextTypeAndDevice) {
  ASSERT_EQ(sharedPools, devices.size() * 2);
 }
 
-INSTANTIATE_TEST_SUITE_P(urUsmPoolManagerTest, urUsmPoolManagerTest,
+INSTANTIATE_TEST_SUITE_P(urUsmPoolDescriptorTest, urUsmPoolDescriptorTest,
  ::testing::Values(nullptr));
 
 // TODO: add test with sub-devices
+
+struct urUsmPoolManagerTest : public uur::urContextTest {
+ void SetUp() override {
+ UUR_RETURN_ON_FATAL_FAILURE(urContextTest::SetUp());
+ auto [ret, descs] = usm::pool_descriptor::create(nullptr, context);
+ ASSERT_EQ(ret, UR_RESULT_SUCCESS);
+ poolDescriptors = descs;
+ }
+
+ std::vector<usm::pool_descriptor> poolDescriptors;
+};
+
+TEST_P(urUsmPoolManagerTest, poolManagerPopulate) {
+ auto [ret, manager] = usm::pool_manager<usm::pool_descriptor>::create();
+ ASSERT_EQ(ret, UR_RESULT_SUCCESS);
+
+ for (auto &desc : poolDescriptors) {
+ // Populate the pool manager
+ auto pool = nullPoolCreate();
+ ASSERT_NE(pool, nullptr);
+ auto poolUnique = umf::pool_unique_handle_t(pool, umfPoolDestroy);
+ ASSERT_NE(poolUnique, nullptr);
+ ret = manager.addPool(desc, poolUnique);
+ ASSERT_EQ(ret, UR_RESULT_SUCCESS);
+ }
+
+ for (auto &desc : poolDescriptors) {
+ // Confirm that there is a pool for each descriptor
+ auto hPoolOpt = manager.getPool(desc);
+ ASSERT_TRUE(hPoolOpt.has_value());
+ ASSERT_NE(hPoolOpt.value(), nullptr);
+ }
+}
+
+TEST_P(urUsmPoolManagerTest, poolManagerInsertExisting) {
+ auto [ret, manager] = usm::pool_manager<usm::pool_descriptor>::create();
+ ASSERT_EQ(ret, UR_RESULT_SUCCESS);
+
+ auto desc = poolDescriptors[0];
+
+ auto pool = nullPoolCreate();
+ ASSERT_NE(pool, nullptr);
+ auto poolUnique = umf::pool_unique_handle_t(pool, umfPoolDestroy);
+ ASSERT_NE(poolUnique, nullptr);
+
+ ret = manager.addPool(desc, poolUnique);
+ ASSERT_EQ(ret, UR_RESULT_SUCCESS);
+
+ // Inserting an existing key should return an error
+ ret = manager.addPool(desc, poolUnique);
+ ASSERT_EQ(ret, UR_RESULT_ERROR_INVALID_ARGUMENT);
+}
+
+TEST_P(urUsmPoolManagerTest, poolManagerGetNonexistant) {
+ auto [ret, manager] = usm::pool_manager<usm::pool_descriptor>::create();
+ ASSERT_EQ(ret, UR_RESULT_SUCCESS);
+
+ for (auto &desc : poolDescriptors) {
+ auto hPool = manager.getPool(desc);
+ ASSERT_FALSE(hPool.has_value());
+ }
+}
+
+UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urUsmPoolManagerTest);