From 868233c78693d9001318985d99f520f018fe2b37 Mon Sep 17 00:00:00 2001
From: Igor Chorazewicz <igor.chorazewicz@intel.com>
Date: Tue, 31 Dec 2024 03:01:04 +0000
Subject: [PATCH] [L0 v2] do not use usm pool free for native handles

as they can be allocated by the user using L0 API directly
(and then UMF will not know about those pointers and umfFree
will do nothing).
---
 source/adapters/level_zero/v2/memory.cpp | 15 ++++++---------
 source/adapters/level_zero/v2/usm.cpp    |  8 +++++++-
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/source/adapters/level_zero/v2/memory.cpp b/source/adapters/level_zero/v2/memory.cpp
index 4de2a588a3..3ec16e8352 100644
--- a/source/adapters/level_zero/v2/memory.cpp
+++ b/source/adapters/level_zero/v2/memory.cpp
@@ -115,10 +115,7 @@ ur_integrated_mem_handle_t::ur_integrated_mem_handle_t(
     if (!ownHostPtr) {
       return;
     }
-    auto ret = hContext->getDefaultUSMPool()->free(ptr);
-    if (ret != UR_RESULT_SUCCESS) {
-      logger::error("Failed to free host memory: {}", ret);
-    }
+    ZE_CALL_NOCHECK(zeMemFree, (hContext->getZeHandle(), ptr));
   });
 }
 
@@ -234,10 +231,7 @@ ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(
           if (!ownZePtr) {
             return;
           }
-          auto ret = hContext->getDefaultUSMPool()->free(ptr);
-          if (ret != UR_RESULT_SUCCESS) {
-            logger::error("Failed to free device memory: {}", ret);
-          }
+          ZE_CALL_NOCHECK(zeMemFree, (hContext->getZeHandle(), ptr));
         });
   }
 }
@@ -310,7 +304,10 @@ void *ur_discrete_mem_handle_t::mapHostPtr(
   usm_unique_ptr_t mappedPtr =
       usm_unique_ptr_t(ptr, [ownsAlloc = bool(mapToPtr), this](void *p) {
         if (ownsAlloc) {
-          UR_CALL_THROWS(hContext->getDefaultUSMPool()->free(p));
+          auto ret = hContext->getDefaultUSMPool()->free(p);
+          if (ret != UR_RESULT_SUCCESS) {
+            logger::error("Failed to mapped memory: {}", ret);
+          }
         }
       });
 
diff --git a/source/adapters/level_zero/v2/usm.cpp b/source/adapters/level_zero/v2/usm.cpp
index 4e9c264182..4f22221562 100644
--- a/source/adapters/level_zero/v2/usm.cpp
+++ b/source/adapters/level_zero/v2/usm.cpp
@@ -231,7 +231,13 @@ ur_result_t ur_usm_pool_handle_t_::allocate(
 }
 
 ur_result_t ur_usm_pool_handle_t_::free(void *ptr) {
-  return umf::umf2urResult(umfFree(ptr));
+  auto umfPool = umfPoolByPtr(ptr);
+  if (umfPool) {
+    return umf::umf2urResult(umfPoolFree(umfPool, ptr));
+  } else {
+    logger::error("Failed to find pool for pointer: {}", ptr);
+    return UR_RESULT_ERROR_INVALID_VALUE;
+  }
 }
 
 namespace ur::level_zero {