[UR][L0] Add initial USM alloc enqueue API

Co-authored-by: Michał Staniewski <[email protected]>

Author: kswiecicki

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

@@ -58,12 +58,14 @@ struct ur_context_handle_t_ : _ur_object {
   ur_context_handle_t_(ze_context_handle_t ZeContext, uint32_t NumDevices,
                        const ur_device_handle_t *Devs, bool OwnZeContext)
       : ZeContext{ZeContext}, Devices{Devs, Devs + NumDevices},
-        NumDevices{NumDevices}, DefaultPool{this, nullptr, !UseUSMAllocator} {
+        NumDevices{NumDevices}, DefaultPool{this, nullptr, !UseUSMAllocator},
+        AsyncPool{this, nullptr} {
     OwnNativeHandle = OwnZeContext;
   }
 
   ur_context_handle_t_(ze_context_handle_t ZeContext)
-      : ZeContext{ZeContext}, DefaultPool{this, nullptr} {}
+      : ZeContext{ZeContext}, DefaultPool{this, nullptr},
+        AsyncPool{this, nullptr} {}
 
   // A L0 context handle is primarily used during creation and management of
   // resources that may be used by multiple devices.
@@ -126,6 +128,9 @@ struct ur_context_handle_t_ : _ur_object {
   // the 'UseUSMAllocator' variable value.
   ur_usm_pool_handle_t_ DefaultPool;
 
+  // USM pools for async allocations.
+  ur_usm_pool_handle_t_ AsyncPool;
+
   // Map associating pools created with urUsmPoolCreate and internal pools
   std::list<ur_usm_pool_handle_t> UsmPoolHandles{};
 

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

@@ -594,124 +594,213 @@ ur_result_t urUSMReleaseExp(ur_context_handle_t Context, void *HostPtr) {
   return UR_RESULT_SUCCESS;
 }
 
+static ur_result_t enqueueUSMAllocHelper(
+    ur_queue_handle_t Queue, ur_usm_pool_handle_t Pool, const size_t Size,
+    const ur_exp_async_usm_alloc_properties_t *Properties,
+    uint32_t NumEventsInWaitList, const ur_event_handle_t *EventWaitList,
+    void **RetMem, ur_event_handle_t *OutEvent, ur_usm_type_t Type) {
+  std::ignore = Pool;
+  std::ignore = Properties;
+
+  std::scoped_lock<ur_shared_mutex> lock(Queue->Mutex);
+
+  bool UseCopyEngine = false;
+  _ur_ze_event_list_t TmpWaitList;
+  UR_CALL(TmpWaitList.createAndRetainUrZeEventList(
+      NumEventsInWaitList, EventWaitList, Queue, UseCopyEngine));
+
+  bool OkToBatch = true;
+  // Get a new command list to be used on this call
+  ur_command_list_ptr_t CommandList{};
+  UR_CALL(Queue->Context->getAvailableCommandList(
+      Queue, CommandList, UseCopyEngine, NumEventsInWaitList, EventWaitList,
+      OkToBatch, nullptr /*ForcedCmdQueue*/));
+
+  ze_event_handle_t ZeEvent = nullptr;
+  ur_event_handle_t InternalEvent{};
+  bool IsInternal = OutEvent == nullptr;
+  ur_event_handle_t *Event = OutEvent ? OutEvent : &InternalEvent;
+
+  ur_command_t CommandType = UR_COMMAND_FORCE_UINT32;
+  switch (Type) {
+  case UR_USM_TYPE_HOST:
+    CommandType = UR_COMMAND_ENQUEUE_USM_HOST_ALLOC_EXP;
+    break;
+  case UR_USM_TYPE_DEVICE:
+    CommandType = UR_COMMAND_ENQUEUE_USM_DEVICE_ALLOC_EXP;
+    break;
+  case UR_USM_TYPE_SHARED:
+    CommandType = UR_COMMAND_ENQUEUE_USM_SHARED_ALLOC_EXP;
+    break;
+  default:
+    logger::error("enqueueUSMAllocHelper: unsupported USM type");
+    throw UR_RESULT_ERROR_UNKNOWN;
+  }
+  UR_CALL(createEventAndAssociateQueue(Queue, Event, CommandType, CommandList,
+                                       IsInternal, false));
+  ZeEvent = (*Event)->ZeEvent;
+  (*Event)->WaitList = TmpWaitList;
+
+  // Allocate USM memory
+  ur_usm_pool_handle_t USMPool = nullptr;
+  if (Pool) {
+    USMPool = Pool;
+  } else {
+    USMPool = &Queue->Context->AsyncPool;
+  }
+
+  auto Device = (Type == UR_USM_TYPE_HOST) ? nullptr : Queue->Device;
+  auto Ret =
+      USMPool->allocate(Queue->Context, Device, nullptr, Type, Size, RetMem);
+  if (Ret) {
+    return Ret;
+  }
+
+  // Signal that USM allocation event was finished
+  ZE2UR_CALL(zeCommandListAppendSignalEvent, (CommandList->first, ZeEvent));
+
+  UR_CALL(Queue->executeCommandList(CommandList, false, OkToBatch));
+
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t urEnqueueUSMDeviceAllocExp(
-    ur_queue_handle_t hQueue, ///< [in] handle of the queue object
-    ur_usm_pool_handle_t
-        pPool,         ///< [in][optional] handle of the USM memory pool
-    const size_t size, ///< [in] minimum size in bytes of the USM memory object
+    ur_queue_handle_t Queue,   ///< [in] handle of the queue object
+    ur_usm_pool_handle_t Pool, ///< [in][optional] USM pool descriptor
+    const size_t Size, ///< [in] minimum size in bytes of the USM memory object
                        ///< to be allocated
     const ur_exp_async_usm_alloc_properties_t
-        *pProperties, ///< [in][optional] pointer to the enqueue asynchronous
-                      ///< USM allocation properties
-    uint32_t numEventsInWaitList, ///< [in] size of the event wait list
+        *Properties, ///< [in][optional] pointer to the enqueue asynchronous
+                     ///< USM allocation properties
+    uint32_t NumEventsInWaitList, ///< [in] size of the event wait list
     const ur_event_handle_t
-        *phEventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
-                          ///< pointer to a list of events that must be complete
-                          ///< before the kernel execution. If nullptr, the
-                          ///< numEventsInWaitList must be 0, indicating no wait
-                          ///< events.
-    void **ppMem,         ///< [out] pointer to USM memory object
-    ur_event_handle_t
-        *phEvent ///< [out][optional] return an event object that identifies the
-                 ///< asynchronous USM device allocation
+        *EventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
+                        ///< pointer to a list of events that must be complete
+                        ///< before the kernel execution. If nullptr, the
+                        ///< numEventsInWaitList must be 0, indicating no wait
+                        ///< events.
+    void **Mem,         ///< [out] pointer to USM memory object
+    ur_event_handle_t *OutEvent ///< [out][optional] return an event object that
+                                ///< identifies the async alloc
 ) {
-  std::ignore = hQueue;
-  std::ignore = pPool;
-  std::ignore = size;
-  std::ignore = pProperties;
-  std::ignore = numEventsInWaitList;
-  std::ignore = phEventWaitList;
-  std::ignore = ppMem;
-  std::ignore = phEvent;
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  return enqueueUSMAllocHelper(Queue, Pool, Size, Properties,
+                               NumEventsInWaitList, EventWaitList, Mem,
+                               OutEvent, UR_USM_TYPE_DEVICE);
 }
 
 ur_result_t urEnqueueUSMSharedAllocExp(
-    ur_queue_handle_t hQueue, ///< [in] handle of the queue object
-    ur_usm_pool_handle_t
-        pPool,         ///< [in][optional] handle of the USM memory pool
-    const size_t size, ///< [in] minimum size in bytes of the USM memory object
+    ur_queue_handle_t Queue,   ///< [in] handle of the queue object
+    ur_usm_pool_handle_t Pool, ///< [in][optional] USM pool descriptor
+    const size_t Size, ///< [in] minimum size in bytes of the USM memory object
                        ///< to be allocated
     const ur_exp_async_usm_alloc_properties_t
-        *pProperties, ///< [in][optional] pointer to the enqueue asynchronous
-                      ///< USM allocation properties
-    uint32_t numEventsInWaitList, ///< [in] size of the event wait list
+        *Properties, ///< [in][optional] pointer to the enqueue asynchronous
+                     ///< USM allocation properties
+    uint32_t NumEventsInWaitList, ///< [in] size of the event wait list
     const ur_event_handle_t
-        *phEventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
-                          ///< pointer to a list of events that must be complete
-                          ///< before the kernel execution. If nullptr, the
-                          ///< numEventsInWaitList must be 0, indicating no wait
-                          ///< events.
-    void **ppMem,         ///< [out] pointer to USM memory object
-    ur_event_handle_t
-        *phEvent ///< [out][optional] return an event object that identifies the
-                 ///< asynchronous USM shared allocation
+        *EventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
+                        ///< pointer to a list of events that must be complete
+                        ///< before the kernel execution. If nullptr, the
+                        ///< numEventsInWaitList must be 0, indicating no wait
+                        ///< events.
+    void **Mem,         ///< [out] pointer to USM memory object
+    ur_event_handle_t *OutEvent ///< [out][optional] return an event object that
+                                ///< identifies the async alloc
 ) {
-  std::ignore = hQueue;
-  std::ignore = pPool;
-  std::ignore = size;
-  std::ignore = pProperties;
-  std::ignore = numEventsInWaitList;
-  std::ignore = phEventWaitList;
-  std::ignore = ppMem;
-  std::ignore = phEvent;
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  return enqueueUSMAllocHelper(Queue, Pool, Size, Properties,
+                               NumEventsInWaitList, EventWaitList, Mem,
+                               OutEvent, UR_USM_TYPE_SHARED);
 }
 
 ur_result_t urEnqueueUSMHostAllocExp(
-    ur_queue_handle_t hQueue, ///< [in] handle of the queue object
-    ur_usm_pool_handle_t
-        pPool,         ///< [in][optional] handle of the USM memory pool
-    const size_t size, ///< [in] minimum size in bytes of the USM memory object
+    ur_queue_handle_t Queue,   ///< [in] handle of the queue object
+    ur_usm_pool_handle_t Pool, ///< [in][optional] handle of the USM memory pool
+    const size_t Size, ///< [in] minimum size in bytes of the USM memory object
                        ///< to be allocated
     const ur_exp_async_usm_alloc_properties_t
-        *pProperties, ///< [in][optional] pointer to the enqueue asynchronous
-                      ///< USM allocation properties
-    uint32_t numEventsInWaitList, ///< [in] size of the event wait list
+        *Properties, ///< [in][optional] pointer to the enqueue asynchronous
+                     ///< USM allocation properties
+    uint32_t NumEventsInWaitList, ///< [in] size of the event wait list
     const ur_event_handle_t
-        *phEventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
-                          ///< pointer to a list of events that must be complete
-                          ///< before the kernel execution. If nullptr, the
-                          ///< numEventsInWaitList must be 0, indicating no wait
-                          ///< events.
-    void **ppMem,         ///< [out] pointer to USM memory object
+        *EventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
+                        ///< pointer to a list of events that must be complete
+                        ///< before the kernel execution. If nullptr, the
+                        ///< numEventsInWaitList must be 0, indicating no wait
+                        ///< events.
+    void **Mem,         ///< [out] pointer to USM memory object
     ur_event_handle_t
-        *phEvent ///< [out][optional] return an event object that identifies the
-                 ///< asynchronous USM host allocation
+        *OutEvent ///< [out][optional] return an event object that identifies
+                  ///< the asynchronous USM device allocation
 ) {
-  std::ignore = hQueue;
-  std::ignore = pPool;
-  std::ignore = size;
-  std::ignore = pProperties;
-  std::ignore = numEventsInWaitList;
-  std::ignore = phEventWaitList;
-  std::ignore = ppMem;
-  std::ignore = phEvent;
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  return enqueueUSMAllocHelper(Queue, Pool, Size, Properties,
+                               NumEventsInWaitList, EventWaitList, Mem,
+                               OutEvent, UR_USM_TYPE_HOST);
 }
 
 ur_result_t urEnqueueUSMFreeExp(
-    ur_queue_handle_t hQueue, ///< [in] handle of the queue object
-    ur_usm_pool_handle_t
-        pPool,  ///< [in][optional] handle of the USM memory pooliptor
-    void *pMem, ///< [in] pointer to USM memory object
-    uint32_t numEventsInWaitList, ///< [in] size of the event wait list
+    ur_queue_handle_t Queue,      ///< [in] handle of the queue object
+    ur_usm_pool_handle_t Pool,    ///< [in][optional] USM pool descriptor
+    void *Mem,                    ///< [in] pointer to USM memory object
+    uint32_t NumEventsInWaitList, ///< [in] size of the event wait list
     const ur_event_handle_t
-        *phEventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
-                          ///< pointer to a list of events that must be complete
-                          ///< before the kernel execution. If nullptr, the
-                          ///< numEventsInWaitList must be 0, indicating no wait
-                          ///< events.
-    ur_event_handle_t *phEvent ///< [out][optional] return an event object that
-                               ///< identifies the asynchronous USM deallocation
+        *EventWaitList, ///< [in][optional][range(0, numEventsInWaitList)]
+                        ///< pointer to a list of events that must be complete
+                        ///< before the kernel execution. If nullptr, the
+                        ///< numEventsInWaitList must be 0, indicating no wait
+                        ///< events.
+    ur_event_handle_t *OutEvent ///< [out][optional] return an event object that
+                                ///< identifies the async alloc
 ) {
-  std::ignore = hQueue;
-  std::ignore = pPool;
-  std::ignore = pMem;
-  std::ignore = numEventsInWaitList;
-  std::ignore = phEventWaitList;
-  std::ignore = phEvent;
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  std::ignore = Pool;
+
+  std::scoped_lock<ur_shared_mutex> lock(Queue->Mutex);
+
+  bool UseCopyEngine = false;
+  _ur_ze_event_list_t TmpWaitList;
+  UR_CALL(TmpWaitList.createAndRetainUrZeEventList(
+      NumEventsInWaitList, EventWaitList, Queue, UseCopyEngine));
+
+  bool OkToBatch = false;
+  // Get a new command list to be used on this call
+  ur_command_list_ptr_t CommandList{};
+  UR_CALL(Queue->Context->getAvailableCommandList(
+      Queue, CommandList, UseCopyEngine, NumEventsInWaitList, EventWaitList,
+      OkToBatch, nullptr /*ForcedCmdQueue*/));
+
+  ze_event_handle_t ZeEvent = nullptr;
+  ur_event_handle_t InternalEvent{};
+  bool IsInternal = OutEvent == nullptr;
+  ur_event_handle_t *Event = OutEvent ? OutEvent : &InternalEvent;
+
+  UR_CALL(createEventAndAssociateQueue(Queue, Event,
+                                       UR_COMMAND_ENQUEUE_USM_FREE_EXP,
+                                       CommandList, IsInternal, false));
+  ZeEvent = (*Event)->ZeEvent;
+  (*Event)->WaitList = TmpWaitList;
+
+  const auto &ZeCommandList = CommandList->first;
+  const auto &WaitList = (*Event)->WaitList;
+  if (WaitList.Length) {
+    ZE2UR_CALL(zeCommandListAppendWaitOnEvents,
+               (ZeCommandList, WaitList.Length, WaitList.ZeEventList));
+
+    // Wait for commands execution until USM can be freed
+    UR_CALL(
+        Queue->executeCommandList(CommandList, true, OkToBatch)); // Blocking
+  }
+
+  // Free USM memory
+  auto Ret = USMFreeHelper(Queue->Context, Mem);
+  if (Ret) {
+    return Ret;
+  }
+
+  // Signal that USM free event was finished
+  ZE2UR_CALL(zeCommandListAppendSignalEvent, (ZeCommandList, ZeEvent));
+
+  UR_CALL(Queue->executeCommandList(CommandList, false, OkToBatch));
+
+  return UR_RESULT_SUCCESS;
 }
 
 ur_result_t UR_APICALL urUSMPoolCreateExp(

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

@@ -40,6 +40,17 @@ function(add_adapter_tests adapter)
 
         add_dependencies(test-adapter-${adapter}
             generate_device_binaries kernel_names_header)
+
+
+        if("${adapter}" STREQUAL "level_zero")
+            add_adapter_test(level_zero_enqueue_alloc
+                FIXTURE KERNELS
+                SOURCES
+                    enqueue_alloc.cpp
+                ENVIRONMENT
+                    "UR_ADAPTERS_FORCE_LOAD=\"$<TARGET_FILE:ur_adapter_level_zero>\""
+            )
+        endif()
     endif()
 
     if(NOT WIN32 AND NOT UR_STATIC_ADAPTER_L0)