[SYCL][UR][L0] Access UMF pool handles through usm::pool_manager

source/adapters/level_zero/context.cpp

@@ -36,7 +36,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urContextCreate(
     ur_context_handle_t_ *Context =
         new ur_context_handle_t_(ZeContext, DeviceCount, Devices, true);
 
-    Context->initialize();
+    auto Ret = Context->initialize();
+    if (Ret) {
+      delete Context;
+      return Ret;
+    }
+
     *RetContext = reinterpret_cast<ur_context_handle_t>(Context);
     if (IndirectAccessTrackingEnabled) {
       std::scoped_lock<ur_shared_mutex> Lock(Platform->ContextsMutex);
@@ -178,111 +183,120 @@ UR_APIEXPORT ur_result_t UR_APICALL urContextSetExtendedDeleter(
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
+// Template helper function for creating USM pools for given pool descriptor.
+template <typename P, typename... Args>
+std::pair<ur_result_t, umf::pool_unique_handle_t>
+createUMFPoolForDesc(usm::pool_descriptor &Desc, Args &&...args) {
+  umf_result_t UmfRet = UMF_RESULT_SUCCESS;
+  umf::provider_unique_handle_t MemProvider = nullptr;
+
+  switch (Desc.type) {
+  case UR_USM_TYPE_HOST: {
+    std::tie(UmfRet, MemProvider) =
+        umf::memoryProviderMakeUnique<L0HostMemoryProvider>(Desc.hContext,
+                                                            Desc.hDevice);
+    break;
+  }
+  case UR_USM_TYPE_DEVICE: {
+    std::tie(UmfRet, MemProvider) =
+        umf::memoryProviderMakeUnique<L0DeviceMemoryProvider>(Desc.hContext,
+                                                              Desc.hDevice);
+    break;
+  }
+  case UR_USM_TYPE_SHARED: {
+    if (Desc.deviceReadOnly) {
+      std::tie(UmfRet, MemProvider) =
+          umf::memoryProviderMakeUnique<L0SharedReadOnlyMemoryProvider>(
+              Desc.hContext, Desc.hDevice);
+    } else {
+      std::tie(UmfRet, MemProvider) =
+          umf::memoryProviderMakeUnique<L0SharedMemoryProvider>(Desc.hContext,
+                                                                Desc.hDevice);
+    }
+    break;
+  }
+  default:
+    UmfRet = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+  }
+
+  if (UmfRet)
+    return std::pair<ur_result_t, umf::pool_unique_handle_t>{
+        umf::umf2urResult(UmfRet), nullptr};
+
+  umf::pool_unique_handle_t Pool = nullptr;
+  std::tie(UmfRet, Pool) =
+      umf::poolMakeUnique<P, 1>({std::move(MemProvider)}, args...);
+
+  return std::pair<ur_result_t, umf::pool_unique_handle_t>{
+      umf::umf2urResult(UmfRet), std::move(Pool)};
+};
+
 ur_result_t ur_context_handle_t_::initialize() {
 
-  // Helper lambda to create various USM allocators for a device.
-  // Note that the CCS devices and their respective subdevices share a
-  // common ze_device_handle and therefore, also share USM allocators.
-  auto createUSMAllocators = [this](ur_device_handle_t Device) {
-    auto MemProvider = umf::memoryProviderMakeUnique<L0DeviceMemoryProvider>(
-                           reinterpret_cast<ur_context_handle_t>(this), Device)
-                           .second;
-    DeviceMemPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(umf::poolMakeUnique<usm::DisjointPool, 1>(
-                            {std::move(MemProvider)},
-                            DisjointPoolConfigInstance
-                                .Configs[usm::DisjointPoolMemType::Device])
-                            .second));
-
-    MemProvider = umf::memoryProviderMakeUnique<L0SharedMemoryProvider>(
-                      reinterpret_cast<ur_context_handle_t>(this), Device)
-                      .second;
-    SharedMemPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(umf::poolMakeUnique<usm::DisjointPool, 1>(
-                            {std::move(MemProvider)},
-                            DisjointPoolConfigInstance
-                                .Configs[usm::DisjointPoolMemType::Shared])
-                            .second));
-
-    MemProvider = umf::memoryProviderMakeUnique<L0SharedReadOnlyMemoryProvider>(
-                      reinterpret_cast<ur_context_handle_t>(this), Device)
-                      .second;
-    SharedReadOnlyMemPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(
-            umf::poolMakeUnique<usm::DisjointPool, 1>(
-                {std::move(MemProvider)},
-                DisjointPoolConfigInstance
-                    .Configs[usm::DisjointPoolMemType::SharedReadOnly])
-                .second));
-
-    MemProvider = umf::memoryProviderMakeUnique<L0DeviceMemoryProvider>(
-                      reinterpret_cast<ur_context_handle_t>(this), Device)
-                      .second;
-    DeviceMemProxyPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(
-            umf::poolMakeUnique<USMProxyPool, 1>({std::move(MemProvider)})
-                .second));
-
-    MemProvider = umf::memoryProviderMakeUnique<L0SharedMemoryProvider>(
-                      reinterpret_cast<ur_context_handle_t>(this), Device)
-                      .second;
-    SharedMemProxyPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(
-            umf::poolMakeUnique<USMProxyPool, 1>({std::move(MemProvider)})
-                .second));
-
-    MemProvider = umf::memoryProviderMakeUnique<L0SharedReadOnlyMemoryProvider>(
-                      reinterpret_cast<ur_context_handle_t>(this), Device)
-                      .second;
-    SharedReadOnlyMemProxyPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(
-            umf::poolMakeUnique<USMProxyPool, 1>({std::move(MemProvider)})
-                .second));
-  };
+  auto Context = reinterpret_cast<ur_context_handle_t>(this);
+  ur_result_t Ret;
 
-  // Recursive helper to call createUSMAllocators for all sub-devices
-  std::function<void(ur_device_handle_t)> createUSMAllocatorsRecursive;
-  createUSMAllocatorsRecursive =
-      [createUSMAllocators,
-       &createUSMAllocatorsRecursive](ur_device_handle_t Device) -> void {
-    createUSMAllocators(Device);
-    for (auto &SubDevice : Device->SubDevices)
-      createUSMAllocatorsRecursive(SubDevice);
-  };
+  // Initialize pool managers.
+  std::tie(Ret, PoolManager) =
+      usm::pool_manager<usm::pool_descriptor>::create();
+  if (Ret) {
+    urPrint("urContextCreate: unexpected internal error\n");
+    return Ret;
+  }
 
-  // Create USM pool for each pair (device, context).
-  //
-  for (auto &Device : Devices) {
-    createUSMAllocatorsRecursive(Device);
+  std::tie(Ret, ProxyPoolManager) =
+      usm::pool_manager<usm::pool_descriptor>::create();
+  if (Ret) {
+    urPrint("urContextCreate: unexpected internal error\n");
+    return Ret;
+  }
+
+  std::vector<usm::pool_descriptor> Descs;
+  // Create pool descriptor for every device and subdevice.
+  std::tie(Ret, Descs) = usm::pool_descriptor::create(nullptr, Context);
+  if (Ret) {
+    urPrint("urContextCreate: unexpected internal error\n");
+    return Ret;
   }
-  // Create USM pool for host. Device and Shared USM allocations
-  // are device-specific. Host allocations are not device-dependent therefore
-  // we don't need a map with device as key.
-  auto MemProvider = umf::memoryProviderMakeUnique<L0HostMemoryProvider>(
-                         reinterpret_cast<ur_context_handle_t>(this), nullptr)
-                         .second;
-  HostMemPool =
-      umf::poolMakeUnique<usm::DisjointPool, 1>(
-          {std::move(MemProvider)},
-          DisjointPoolConfigInstance.Configs[usm::DisjointPoolMemType::Host])
-          .second;
-
-  MemProvider = umf::memoryProviderMakeUnique<L0HostMemoryProvider>(
-                    reinterpret_cast<ur_context_handle_t>(this), nullptr)
-                    .second;
-  HostMemProxyPool =
-      umf::poolMakeUnique<USMProxyPool, 1>({std::move(MemProvider)}).second;
-
-  // We may allocate memory to this root device so create allocators.
-  if (SingleRootDevice &&
-      DeviceMemPools.find(SingleRootDevice->ZeDevice) == DeviceMemPools.end()) {
-    createUSMAllocators(SingleRootDevice);
+
+  auto descTypeToDisjointPoolType =
+      [](usm::pool_descriptor &Desc) -> usm::DisjointPoolMemType {
+    switch (Desc.type) {
+    case UR_USM_TYPE_HOST:
+      return usm::DisjointPoolMemType::Host;
+    case UR_USM_TYPE_DEVICE:
+      return usm::DisjointPoolMemType::Device;
+    case UR_USM_TYPE_SHARED:
+      return (Desc.deviceReadOnly) ? usm::DisjointPoolMemType::SharedReadOnly
+                                   : usm::DisjointPoolMemType::Shared;
+    default:
+      // Should not be reached.
+      ur::unreachable();
+    }
+  };
+
+  // Create USM pool for each pool descriptor and add it to pool manager.
+  for (auto &Desc : Descs) {
+    umf::pool_unique_handle_t Pool = nullptr;
+    auto PoolType = descTypeToDisjointPoolType(Desc);
+
+    std::tie(Ret, Pool) = createUMFPoolForDesc<usm::DisjointPool>(
+        Desc, DisjointPoolConfigInstance.Configs[PoolType]);
+    if (Ret) {
+      urPrint("urContextCreate: unexpected internal error\n");
+      return Ret;
+    }
+
+    PoolManager.addPool(Desc, Pool);
+
+    umf::pool_unique_handle_t ProxyPool = nullptr;
+    std::tie(Ret, ProxyPool) = createUMFPoolForDesc<USMProxyPool>(Desc);
+    if (Ret) {
+      urPrint("urContextCreate: unexpected internal error\n");
+      return Ret;
+    }
+
+    ProxyPoolManager.addPool(Desc, ProxyPool);
   }
 
   // Create the immediate command list to be used for initializations.

source/adapters/level_zero/context.hpp

@@ -26,6 +26,7 @@
 #include "queue.hpp"
 
 #include <umf_helpers.hpp>
+#include <ur_pool_manager.hpp>
 
 struct ur_context_handle_t_ : _ur_object {
   ur_context_handle_t_(ze_context_handle_t ZeContext, uint32_t NumDevices,
@@ -96,15 +97,8 @@ struct ur_context_handle_t_ : _ur_object {
 
   // Store USM pool for USM shared and device allocations. There is 1 memory
   // pool per each pair of (context, device) per each memory type.
-  std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
-      DeviceMemPools;
-  std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
-      SharedMemPools;
-  std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
-      SharedReadOnlyMemPools;
-
-  // Store the host memory pool. It does not depend on any device.
-  umf::pool_unique_handle_t HostMemPool;
+  usm::pool_manager<usm::pool_descriptor> PoolManager;
+  usm::pool_manager<usm::pool_descriptor> ProxyPoolManager;
 
   // Allocation-tracking proxy pools for direct allocations. No pooling used.
   std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
@@ -252,3 +246,8 @@ struct ur_context_handle_t_ : _ur_object {
 // mutex guarding the container with contexts because the context can be removed
 // from the list of tracked contexts.
 ur_result_t ContextReleaseHelper(ur_context_handle_t Context);
+
+// Template helper function for creating USM pools for given pool descriptor.
+template <typename P, typename... Args>
+std::pair<ur_result_t, umf::pool_unique_handle_t>
+createUMFPoolForDesc(usm::pool_descriptor &Desc, Args &&...args);

source/adapters/level_zero/device.cpp

@@ -683,7 +683,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(
         }
       }
     }
-    return ReturnValue(std::min(GlobalMemSize, FreeMemory));
+    return ReturnValue((std::min)(GlobalMemSize, FreeMemory));
   }
   case UR_DEVICE_INFO_MEMORY_CLOCK_RATE: {
     // If there are not any memory modules then return 0.