[UR][L0 v2] Port USM alloc to adapter v2

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

@@ -1,6 +1,9 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
+// UNSUPPORTED: level_zero_v2_adapter
+// UNSUPPORTED-TRACKER: https://github.com/intel/llvm/issues/18488
+
 #include <iostream>
 #include <sycl/detail/core.hpp>
 #include <sycl/properties/queue_properties.hpp>

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

@@ -1,6 +1,9 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
+// UNSUPPORTED: level_zero_v2_adapter
+// UNSUPPORTED-TRACKER: https://github.com/intel/llvm/issues/18488
+
 #include <iostream>
 #include <sycl/detail/core.hpp>
 #include <sycl/properties/queue_properties.hpp>

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

@@ -1,6 +1,9 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
+// UNSUPPORTED: level_zero_v2_adapter
+// UNSUPPORTED-TRACKER: https://github.com/intel/llvm/issues/18488
+
 #include <iostream>
 #include <sycl/detail/core.hpp>
 #include <sycl/properties/queue_properties.hpp>

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

@@ -1,6 +1,9 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
+// UNSUPPORTED: level_zero_v2_adapter
+// UNSUPPORTED-TRACKER: https://github.com/intel/llvm/issues/18488
+
 #include <iostream>
 #include <sycl/detail/core.hpp>
 #include <sycl/properties/queue_properties.hpp>

sycl/test-e2e/AsyncAlloc/lit.local.cfg

@@ -1,3 +1 @@
 config.required_features += ['aspect-ext_oneapi_async_memory_alloc']
-# V2 adapter does not support async alloc api yet
-config.unsupported_features += ['level_zero_v2_adapter']

unified-runtime/source/adapters/level_zero/CMakeLists.txt

@@ -141,6 +141,7 @@ if(UR_BUILD_ADAPTER_L0_V2)
         ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/command_buffer_command.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/enqueued_pool.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/image_common.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/ur_interface_loader.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/platform.cpp

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

@@ -93,7 +93,7 @@ static ur_result_t enqueueUSMAllocHelper(
     break;
   default:
     UR_LOG(ERR, "enqueueUSMAllocHelper: unsupported USM type");
-    throw UR_RESULT_ERROR_UNKNOWN;
+    throw UR_RESULT_ERROR_INVALID_ARGUMENT;
   }
   UR_CALL(createEventAndAssociateQueue(Queue, Event, CommandType, CommandList,
                                        IsInternal, false));
@@ -247,6 +247,7 @@ ur_result_t urEnqueueUSMFreeExp(
   }
 
   size_t size = umfPoolMallocUsableSize(hPool, Mem);
+  (*Event)->RefCount.increment();
   usmPool->AsyncPool.insert(Mem, size, *Event, Queue);
 
   // Signal that USM free event was finished

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

@@ -9,15 +9,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "enqueued_pool.hpp"
-#include "event.hpp"
 
 #include <ur_api.h>
 
 EnqueuedPool::~EnqueuedPool() { cleanup(); }
 
 std::optional<EnqueuedPool::Allocation>
-EnqueuedPool::getBestFit(size_t Size, size_t Alignment,
-                         ur_queue_handle_t Queue) {
+EnqueuedPool::getBestFit(size_t Size, size_t Alignment, void *Queue) {
   auto Lock = std::lock_guard(Mutex);
 
   Allocation Alloc = {nullptr, Size, nullptr, Queue, Alignment};
@@ -47,12 +45,11 @@ EnqueuedPool::getBestFit(size_t Size, size_t Alignment,
 }
 
 void EnqueuedPool::insert(void *Ptr, size_t Size, ur_event_handle_t Event,
-                          ur_queue_handle_t Queue) {
+                          void *Queue) {
   auto Lock = std::lock_guard(Mutex);
 
   uintptr_t Address = (uintptr_t)Ptr;
   size_t Alignment = Address & (~Address + 1);
-  Event->RefCount.increment();
 
   Freelist.emplace(Allocation{Ptr, Size, Event, Queue, Alignment});
 }
@@ -67,14 +64,15 @@ bool EnqueuedPool::cleanup() {
     auto umfRet [[maybe_unused]] = umfPoolFree(hPool, It.Ptr);
     assert(umfRet == UMF_RESULT_SUCCESS);
 
-    urEventReleaseInternal(It.Event);
+    if (It.Event)
+      eventRelease(It.Event);
   }
   Freelist.clear();
 
   return FreedAllocations;
 }
 
-bool EnqueuedPool::cleanupForQueue(ur_queue_handle_t Queue) {
+bool EnqueuedPool::cleanupForQueue(void *Queue) {
   auto Lock = std::lock_guard(Mutex);
 
   Allocation Alloc = {nullptr, 0, nullptr, Queue, 0};
@@ -90,7 +88,8 @@ bool EnqueuedPool::cleanupForQueue(ur_queue_handle_t Queue) {
     auto umfRet [[maybe_unused]] = umfPoolFree(hPool, It->Ptr);
     assert(umfRet == UMF_RESULT_SUCCESS);
 
-    urEventReleaseInternal(It->Event);
+    if (It->Event)
+      eventRelease(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

@@ -22,17 +22,24 @@ class EnqueuedPool {
     void *Ptr;
     size_t Size;
     ur_event_handle_t Event;
-    ur_queue_handle_t Queue;
+    // Queue handle, used as an identifier for the associated queue.
+    // This can either be a `ur_queue_handle_t` or a pointer to a v2 queue
+    // object.
+    void *Queue;
     size_t Alignment;
   };
 
+  using event_release_callback_t = ur_result_t (*)(ur_event_handle_t);
+
+  EnqueuedPool(event_release_callback_t eventRelease)
+      : eventRelease(eventRelease) {}
+
   ~EnqueuedPool();
   std::optional<Allocation> getBestFit(size_t Size, size_t Alignment,
-                                       ur_queue_handle_t Queue);
-  void insert(void *Ptr, size_t Size, ur_event_handle_t Event,
-              ur_queue_handle_t Queue);
+                                       void *Queue);
+  void insert(void *Ptr, size_t Size, ur_event_handle_t Event, void *Queue);
   bool cleanup();
-  bool cleanupForQueue(ur_queue_handle_t Queue);
+  bool cleanupForQueue(void *Queue);
 
 private:
   struct Comparator {
@@ -53,4 +60,5 @@ class EnqueuedPool {
   using AllocationSet = std::set<Allocation, Comparator>;
   ur_mutex Mutex;
   AllocationSet Freelist;
+  event_release_callback_t eventRelease;
 };

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

@@ -12,6 +12,7 @@
 #include "common.hpp"
 
 #include "enqueued_pool.hpp"
+#include "event.hpp"
 #include "ur_api.h"
 #include "ur_pool_manager.hpp"
 #include <set>
@@ -20,7 +21,10 @@
 usm::DisjointPoolAllConfigs InitializeDisjointPoolConfig();
 
 struct UsmPool {
-  UsmPool(umf::pool_unique_handle_t Pool) : UmfPool(std::move(Pool)) {}
+  UsmPool(umf::pool_unique_handle_t Pool)
+      : UmfPool(std::move(Pool)), AsyncPool([](ur_event_handle_t Event) {
+          return urEventReleaseInternal(Event);
+        }) {}
   umf::pool_unique_handle_t UmfPool;
   // 'AsyncPool' needs to be declared after 'UmfPool' so its destructor is
   // invoked first.

unified-runtime/source/adapters/level_zero/v2/context.cpp

@@ -68,7 +68,7 @@ ur_context_handle_t_::ur_context_handle_t_(ze_context_handle_t hContext,
                                  v2::EVENT_FLAGS_PROFILING_ENABLED)),
       p2pAccessDevices(populateP2PDevices(
           phDevices[0]->Platform->getNumDevices(), this->hDevices)),
-      defaultUSMPool(this, nullptr) {}
+      defaultUSMPool(this, nullptr), asyncPool(this, nullptr) {}
 
 ur_result_t ur_context_handle_t_::retain() {
   RefCount.increment();
@@ -105,6 +105,8 @@ ur_usm_pool_handle_t ur_context_handle_t_::getDefaultUSMPool() {
   return &defaultUSMPool;
 }
 
+ur_usm_pool_handle_t ur_context_handle_t_::getAsyncPool() { return &asyncPool; }
+
 const std::vector<ur_device_handle_t> &
 ur_context_handle_t_::getP2PDevices(ur_device_handle_t hDevice) const {
   return p2pAccessDevices[hDevice->Id.value()];

unified-runtime/source/adapters/level_zero/v2/context.hpp

@@ -29,6 +29,7 @@ struct ur_context_handle_t_ : ur_object {
 
   const std::vector<ur_device_handle_t> &getDevices() const;
   ur_usm_pool_handle_t getDefaultUSMPool();
+  ur_usm_pool_handle_t getAsyncPool();
 
   const std::vector<ur_device_handle_t> &
   getP2PDevices(ur_device_handle_t hDevice) const;
@@ -55,4 +56,5 @@ struct ur_context_handle_t_ : ur_object {
   const std::vector<std::vector<ur_device_handle_t>> p2pAccessDevices;
 
   ur_usm_pool_handle_t_ defaultUSMPool;
+  ur_usm_pool_handle_t_ asyncPool;
 };