[UR][L0] Add pool statistics (#19360)

Support pool query via urUSMPoolGetInfoExp for information about
reserved or in-use memory and their peak sizes.

Fixes intel/llvm#17772.

Author: kswiecicki

sycl/test-e2e/AsyncAlloc/device/memory_pool.cpp

@@ -1,8 +1,8 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
-// XFAIL: level_zero
-// XFAIL-TRACKER: https://github.com/intel/llvm/issues/17772
+// UNSUPPORTED: level_zero_v2_adapter
+// UNSUPPORTED-INTENDED: v2 adapter does not support pool statistics.
 
 #include <iostream>
 #include <sycl/detail/core.hpp>

unified-runtime/source/adapters/level_zero/async_alloc.cpp

@@ -28,22 +28,22 @@ static ur_result_t enqueueUSMAllocHelper(
   std::scoped_lock<ur_shared_mutex> lock(Queue->Mutex);
 
   // Allocate USM memory
-  ur_usm_pool_handle_t USMPool = nullptr;
+  ur_usm_pool_handle_t UrPool = nullptr;
   if (Pool) {
-    USMPool = Pool;
+    UrPool = Pool;
   } else {
-    USMPool = &Queue->Context->AsyncPool;
+    UrPool = &Queue->Context->AsyncPool;
   }
 
   auto Device = (Type == UR_USM_TYPE_HOST) ? nullptr : Queue->Device;
 
   std::vector<ur_event_handle_t> ExtEventWaitList;
   ur_event_handle_t OriginAllocEvent = nullptr;
   auto AsyncAlloc =
-      USMPool->allocateEnqueued(Queue, Device, nullptr, Type, Size);
+      UrPool->allocateEnqueued(Queue, Device, nullptr, Type, Size);
   if (!AsyncAlloc) {
     auto Ret =
-        USMPool->allocate(Queue->Context, Device, nullptr, Type, Size, RetMem);
+        UrPool->allocate(Queue->Context, Device, nullptr, Type, Size, RetMem);
     if (Ret) {
       return Ret;
     }
@@ -235,26 +235,26 @@ ur_result_t urEnqueueUSMFreeExp(
                (ZeCommandList, WaitList.Length, WaitList.ZeEventList));
   }
 
-  umf_memory_pool_handle_t hPool = nullptr;
-  auto umfRet = umfPoolByPtr(Mem, &hPool);
-  if (umfRet != UMF_RESULT_SUCCESS || !hPool) {
+  umf_memory_pool_handle_t UmfPool = nullptr;
+  auto UmfRet = umfPoolByPtr(Mem, &UmfPool);
+  if (UmfRet != UMF_RESULT_SUCCESS || !UmfPool) {
     return USMFreeHelper(Queue->Context, Mem);
   }
 
-  UsmPool *usmPool = nullptr;
-  umfRet = umfPoolGetTag(hPool, (void **)&usmPool);
-  if (umfRet != UMF_RESULT_SUCCESS || usmPool == nullptr) {
+  UsmPool *UsmPool = nullptr;
+  UmfRet = umfPoolGetTag(UmfPool, (void **)&UsmPool);
+  if (UmfRet != UMF_RESULT_SUCCESS || UsmPool == nullptr) {
     return USMFreeHelper(Queue->Context, Mem);
   }
 
-  size_t size = 0;
-  umfRet = umfPoolMallocUsableSize(hPool, Mem, &size);
-  if (umfRet != UMF_RESULT_SUCCESS) {
+  size_t Size = 0;
+  UmfRet = umfPoolMallocUsableSize(UmfPool, Mem, &Size);
+  if (UmfRet != UMF_RESULT_SUCCESS) {
     return USMFreeHelper(Queue->Context, Mem);
   }
 
   (*Event)->RefCount.retain();
-  usmPool->AsyncPool.insert(Mem, size, *Event, Queue);
+  UsmPool->AsyncPool.insert(Mem, Size, *Event, Queue);
 
   // Signal that USM free event was finished
   ZE2UR_CALL(zeCommandListAppendSignalEvent, (ZeCommandList, ZeEvent));

unified-runtime/source/adapters/level_zero/enqueued_pool.cpp

@@ -9,6 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "enqueued_pool.hpp"
+#include "usm.hpp"
 
 #include <ur_api.h>
 
@@ -58,17 +59,13 @@ bool EnqueuedPool::cleanup() {
   auto Lock = std::lock_guard(Mutex);
   auto FreedAllocations = !Freelist.empty();
 
-  auto umfRet [[maybe_unused]] = UMF_RESULT_SUCCESS;
+  auto Ret [[maybe_unused]] = UR_RESULT_SUCCESS;
   for (auto It : Freelist) {
-    umf_memory_pool_handle_t hPool = nullptr;
-    umfRet = umfPoolByPtr(It.Ptr, &hPool);
-    assert(hPool != nullptr);
-
-    umfRet = umfPoolFree(hPool, It.Ptr);
-    assert(umfRet == UMF_RESULT_SUCCESS);
+    Ret = MemFreeFn(It.Ptr);
+    assert(Ret == UR_RESULT_SUCCESS);
 
     if (It.Event)
-      eventRelease(It.Event);
+      EventReleaseFn(It.Event);
   }
   Freelist.clear();
 
@@ -84,17 +81,13 @@ bool EnqueuedPool::cleanupForQueue(void *Queue) {
 
   bool FreedAllocations = false;
 
-  auto umfRet [[maybe_unused]] = UMF_RESULT_SUCCESS;
+  auto Ret [[maybe_unused]] = UR_RESULT_SUCCESS;
   while (It != Freelist.end() && It->Queue == Queue) {
-    umf_memory_pool_handle_t hPool = nullptr;
-    umfRet = umfPoolByPtr(It->Ptr, &hPool);
-    assert(hPool != nullptr);
-
-    umfRet = umfPoolFree(hPool, It->Ptr);
-    assert(umfRet == UMF_RESULT_SUCCESS);
+    Ret = MemFreeFn(It->Ptr);
+    assert(Ret == UR_RESULT_SUCCESS);
 
     if (It->Event)
-      eventRelease(It->Event);
+      EventReleaseFn(It->Event);
 
     // Erase the current allocation and move to the next one
     It = Freelist.erase(It);

unified-runtime/source/adapters/level_zero/enqueued_pool.hpp

@@ -30,9 +30,11 @@ class EnqueuedPool {
   };
 
   using event_release_callback_t = ur_result_t (*)(ur_event_handle_t);
+  using memory_free_callback_t = std::function<ur_result_t(void *)>;
 
-  EnqueuedPool(event_release_callback_t eventRelease)
-      : eventRelease(eventRelease) {}
+  EnqueuedPool(event_release_callback_t EventReleaseFn,
+               memory_free_callback_t MemFreeFn)
+      : EventReleaseFn(EventReleaseFn), MemFreeFn(MemFreeFn) {}
 
   ~EnqueuedPool();
   std::optional<Allocation> getBestFit(size_t Size, size_t Alignment,
@@ -60,5 +62,6 @@ class EnqueuedPool {
   using AllocationSet = std::set<Allocation, Comparator>;
   ur_mutex Mutex;
   AllocationSet Freelist;
-  event_release_callback_t eventRelease;
+  event_release_callback_t EventReleaseFn;
+  memory_free_callback_t MemFreeFn;
 };