[SYCL][UR][L0] Use usm::pool_manager functionality

... to get the UR_USM_ALLOC_INFO_POOL alloc info.
oneapi-src · Jan 23, 2024 · dfa711f · dfa711f
1 parent 137fa18
commit dfa711f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 22 deletions.
diff --git a/source/adapters/level_zero/usm.cpp b/source/adapters/level_zero/usm.cpp
@@ -615,28 +615,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMGetMemAllocInfo(
 
     std::shared_lock<ur_shared_mutex> ContextLock(Context->Mutex);
 
-    auto SearchMatchingPool =
-        [](std::unordered_map<ur_device_handle_t, umf::pool_unique_handle_t>
-               &PoolMap,
-           umf_memory_pool_handle_t UMFPool) {
-          for (auto &PoolPair : PoolMap) {
-            if (PoolPair.second.get() == UMFPool) {
-              return true;
-            }
-          }
-          return false;
-        };
-
     for (auto &Pool : Context->UsmPoolHandles) {
-      if (SearchMatchingPool(Pool->DeviceMemPools, UMFPool)) {
-        return ReturnValue(Pool);
-      }
-      if (SearchMatchingPool(Pool->SharedMemPools, UMFPool)) {
+      if (Pool->PoolManager.hasPool(UMFPool))
         return ReturnValue(Pool);
-      }
-      if (Pool->HostMemPool.get() == UMFPool) {
-        return ReturnValue(Pool);
-      }
     }
 
     return UR_RESULT_ERROR_INVALID_VALUE;

diff --git a/source/common/ur_pool_manager.hpp b/source/common/ur_pool_manager.hpp
@@ -249,13 +249,19 @@ template <typename D> struct pool_manager {
     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);
+            logger::error(
+                "Pool descriptor: {}, doesn't match any existing pool", desc);
             return std::nullopt;
         }
 
         return it->second.get();
     }
+
+    bool hasPool(umf_memory_pool_handle_t hPool) noexcept {
+        return std::any_of(
+            descToPoolMap.begin(), descToPoolMap.end(),
+            [&hPool](const auto &pair) { return hPool == pair.second.get(); });
+    }
 };
 
 } // namespace usm