[L0] Manage UMF pools through usm::pool_manager

kswiecicki · Jan 7, 2025 · e22b17a · e22b17a
1 parent a3d36c1
commit e22b17a
Show file tree

Hide file tree

Showing 5 changed files with 213 additions and 432 deletions.
diff --git a/source/adapters/level_zero/context.cpp b/source/adapters/level_zero/context.cpp
@@ -191,108 +191,6 @@ ur_result_t urContextSetExtendedDeleter(
 
 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;
-    auto UmfDeviceParamsHandle = getUmfParamsHandle(
-        DisjointPoolConfigInstance.Configs[usm::DisjointPoolMemType::Device]);
-    DeviceMemPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(umf::poolMakeUniqueFromOps(umfDisjointPoolOps(),
-                                                   std::move(MemProvider),
-                                                   UmfDeviceParamsHandle.get())
-                            .second));
-
-    MemProvider = umf::memoryProviderMakeUnique<L0SharedMemoryProvider>(
-                      reinterpret_cast<ur_context_handle_t>(this), Device)
-                      .second;
-
-    auto UmfSharedParamsHandle = getUmfParamsHandle(
-        DisjointPoolConfigInstance.Configs[usm::DisjointPoolMemType::Shared]);
-    SharedMemPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(umf::poolMakeUniqueFromOps(umfDisjointPoolOps(),
-                                                   std::move(MemProvider),
-                                                   UmfSharedParamsHandle.get())
-                            .second));
-
-    MemProvider = umf::memoryProviderMakeUnique<L0SharedReadOnlyMemoryProvider>(
-                      reinterpret_cast<ur_context_handle_t>(this), Device)
-                      .second;
-
-    auto UmfSharedROParamsHandle = getUmfParamsHandle(
-        DisjointPoolConfigInstance
-            .Configs[usm::DisjointPoolMemType::SharedReadOnly]);
-    SharedReadOnlyMemPools.emplace(
-        std::piecewise_construct, std::make_tuple(Device->ZeDevice),
-        std::make_tuple(umf::poolMakeUniqueFromOps(
-                            umfDisjointPoolOps(), std::move(MemProvider),
-                            UmfSharedROParamsHandle.get())
-                            .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>(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>(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>(std::move(MemProvider)).second));
-  };
-
-  // 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);
-  };
-
-  // Create USM pool for each pair (device, context).
-  //
-  for (auto &Device : Devices) {
-    createUSMAllocatorsRecursive(Device);
-  }
-  // 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;
-  auto UmfHostParamsHandle = getUmfParamsHandle(
-      DisjointPoolConfigInstance.Configs[usm::DisjointPoolMemType::Host]);
-  HostMemPool =
-      umf::poolMakeUniqueFromOps(umfDisjointPoolOps(), std::move(MemProvider),
-                                 UmfHostParamsHandle.get())
-          .second;
-
-  MemProvider = umf::memoryProviderMakeUnique<L0HostMemoryProvider>(
-                    reinterpret_cast<ur_context_handle_t>(this), nullptr)
-                    .second;
-  HostMemProxyPool =
-      umf::poolMakeUnique<USMProxyPool>(std::move(MemProvider)).second;
-
   // We may allocate memory to this root device so create allocators.
   if (SingleRootDevice &&
       DeviceMemPools.find(SingleRootDevice->ZeDevice) == DeviceMemPools.end()) {

diff --git a/source/adapters/level_zero/context.hpp b/source/adapters/level_zero/context.hpp
@@ -24,6 +24,7 @@
 
 #include "common.hpp"
 #include "queue.hpp"
+#include "usm.hpp"
 
 #include <umf_helpers.hpp>
 
@@ -55,11 +56,14 @@ struct ur_context_handle_t_ : _ur_object {
   ur_context_handle_t_(ze_context_handle_t ZeContext, uint32_t NumDevices,
                        const ur_device_handle_t *Devs, bool OwnZeContext)
       : ZeContext{ZeContext}, Devices{Devs, Devs + NumDevices},
-        NumDevices{NumDevices} {
+        NumDevices{NumDevices}, DefaultPool{this, nullptr}, ProxyPool{this,
+                                                                      nullptr} {
     OwnNativeHandle = OwnZeContext;
   }
 
-  ur_context_handle_t_(ze_context_handle_t ZeContext) : ZeContext{ZeContext} {}
+  ur_context_handle_t_(ze_context_handle_t ZeContext)
+      : ZeContext{ZeContext}, DefaultPool{this, nullptr}, ProxyPool{this,
+                                                                    nullptr} {}
 
   // A L0 context handle is primarily used during creation and management of
   // resources that may be used by multiple devices.
@@ -123,24 +127,10 @@ 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;
+  ur_usm_pool_handle_t_ DefaultPool;
 
   // Allocation-tracking proxy pools for direct allocations. No pooling used.
-  std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
-      DeviceMemProxyPools;
-  std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
-      SharedMemProxyPools;
-  std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
-      SharedReadOnlyMemProxyPools;
-  umf::pool_unique_handle_t HostMemProxyPool;
+  ur_usm_pool_handle_t_ ProxyPool;
 
   // Map associating pools created with urUsmPoolCreate and internal pools
   std::list<ur_usm_pool_handle_t> UsmPoolHandles{};