[SYCL][UR][L0] Replace memory type look-up with UMF tracking (#10807)

This also makes SharedReadOnly allocations tracking obsolete.

Author: kswiecicki

sycl/plugins/unified_runtime/CMakeLists.txt

@@ -6,6 +6,7 @@ if (NOT DEFINED UNIFIED_RUNTIME_LIBRARY OR NOT DEFINED UNIFIED_RUNTIME_INCLUDE_D
   set(UNIFIED_RUNTIME_REPO "https://github.com/oneapi-src/unified-runtime.git")
   set(UNIFIED_RUNTIME_TAG b3cc9ae3f99ca7faff1ba765dd36652fef2cfddd)
 
+  set(UMF_ENABLE_POOL_TRACKING ON)
   message(STATUS "Will fetch Unified Runtime from ${UNIFIED_RUNTIME_REPO}")
   FetchContent_Declare(unified-runtime
     GIT_REPOSITORY    ${UNIFIED_RUNTIME_REPO}
@@ -46,7 +47,6 @@ if (NOT DEFINED UNIFIED_RUNTIME_LIBRARY OR NOT DEFINED UNIFIED_RUNTIME_INCLUDE_D
   set(UNIFIED_RUNTIME_INCLUDE_DIR "${UNIFIED_RUNTIME_SOURCE_DIR}/include")
 endif()
 
-
 add_library(UnifiedRuntime-Headers INTERFACE)
 
 target_include_directories(UnifiedRuntime-Headers

sycl/plugins/unified_runtime/ur/adapters/level_zero/context.cpp

@@ -182,7 +182,7 @@ ur_result_t ur_context_handle_t_::initialize() {
   // 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<USMDeviceMemoryProvider>(
+    auto MemProvider = umf::memoryProviderMakeUnique<L0DeviceMemoryProvider>(
                            reinterpret_cast<ur_context_handle_t>(this), Device)
                            .second;
     DeviceMemPools.emplace(
@@ -193,7 +193,7 @@ ur_result_t ur_context_handle_t_::initialize() {
                                 .Configs[usm::DisjointPoolMemType::Device])
                             .second));
 
-    MemProvider = umf::memoryProviderMakeUnique<USMSharedMemoryProvider>(
+    MemProvider = umf::memoryProviderMakeUnique<L0SharedMemoryProvider>(
                       reinterpret_cast<ur_context_handle_t>(this), Device)
                       .second;
     SharedMemPools.emplace(
@@ -204,10 +204,9 @@ ur_result_t ur_context_handle_t_::initialize() {
                                 .Configs[usm::DisjointPoolMemType::Shared])
                             .second));
 
-    MemProvider =
-        umf::memoryProviderMakeUnique<USMSharedReadOnlyMemoryProvider>(
-            reinterpret_cast<ur_context_handle_t>(this), Device)
-            .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(
@@ -216,6 +215,33 @@ ur_result_t ur_context_handle_t_::initialize() {
                 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));
   };
 
   // Recursive helper to call createUSMAllocators for all sub-devices
@@ -236,7 +262,7 @@ ur_result_t ur_context_handle_t_::initialize() {
   // 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<USMHostMemoryProvider>(
+  auto MemProvider = umf::memoryProviderMakeUnique<L0HostMemoryProvider>(
                          reinterpret_cast<ur_context_handle_t>(this), nullptr)
                          .second;
   HostMemPool =
@@ -245,6 +271,12 @@ ur_result_t ur_context_handle_t_::initialize() {
           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()) {

sycl/plugins/unified_runtime/ur/adapters/level_zero/context.hpp

@@ -101,14 +101,18 @@ struct ur_context_handle_t_ : _ur_object {
   std::unordered_map<ze_device_handle_t, umf::pool_unique_handle_t>
       SharedReadOnlyMemPools;
 
-  // Since L0 native runtime does not distinguisg "shared device_read_only"
-  // vs regular "shared" allocations, we have keep track of it to use
-  // proper memory pool when freeing allocations.
-  std::unordered_set<void *> SharedReadOnlyAllocs;
-
   // Store the host memory pool. It does not depend on any device.
   umf::pool_unique_handle_t HostMemPool;
 
+  // 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;
+
   // We need to store all memory allocations in the context because there could
   // be kernels with indirect access. Kernels with indirect access start to
   // reference all existing memory allocations at the time when they are