Merge pull request #2241 from igchor/native_2

[L0 v2] implement native APIs and urMemBufferPartition

Author: pbalcer

source/adapters/level_zero/helpers/memory_helpers.cpp

@@ -11,20 +11,27 @@
 #include "memory_helpers.hpp"
 #include "../common.hpp"
 
-ze_memory_type_t getMemoryType(ze_context_handle_t hContext, void *ptr) {
+ur_result_t
+getMemoryAttrs(ze_context_handle_t hContext, void *ptr,
+               ze_device_handle_t *hDevice,
+               ZeStruct<ze_memory_allocation_properties_t> *properties) {
   // TODO: use UMF once
   // https://github.com/oneapi-src/unified-memory-framework/issues/687 is
   // implemented
-  ZeStruct<ze_memory_allocation_properties_t> zeMemoryAllocationProperties;
-  ZE2UR_CALL_THROWS(zeMemGetAllocProperties,
-                    (hContext, ptr, &zeMemoryAllocationProperties, nullptr));
-  return zeMemoryAllocationProperties.type;
+  ZE2UR_CALL(zeMemGetAllocProperties, (hContext, ptr, properties, hDevice));
+  return UR_RESULT_SUCCESS;
 }
 
 bool maybeImportUSM(ze_driver_handle_t hTranslatedDriver,
                     ze_context_handle_t hContext, void *ptr, size_t size) {
-  if (ZeUSMImport.Enabled && ptr != nullptr &&
-      getMemoryType(hContext, ptr) == ZE_MEMORY_TYPE_UNKNOWN) {
+  if (!ZeUSMImport.Enabled || ptr == nullptr) {
+    return false;
+  }
+
+  ZeStruct<ze_memory_allocation_properties_t> properties;
+  auto ret = getMemoryAttrs(hContext, ptr, nullptr, &properties);
+
+  if (ret == UR_RESULT_SUCCESS && properties.type == ZE_MEMORY_TYPE_UNKNOWN) {
     // Promote the host ptr to USM host memory
     ZeUSMImport.doZeUSMImport(hTranslatedDriver, ptr, size);
     return true;

source/adapters/level_zero/helpers/memory_helpers.hpp

@@ -12,6 +12,10 @@
 #include <ur_api.h>
 #include <ze_api.h>
 
+#include <utility>
+
+#include "../common.hpp"
+
 // If USM Import feature is enabled and hostptr is supplied,
 // import the hostptr if not already imported into USM.
 // Data transfer rate is maximized when both source and destination
@@ -20,7 +24,11 @@
 bool maybeImportUSM(ze_driver_handle_t hTranslatedDriver,
                     ze_context_handle_t hContext, void *ptr, size_t size);
 
-ze_memory_type_t getMemoryType(ze_context_handle_t hContext, void *ptr);
+// Get memory attributes for a given pointer
+ur_result_t
+getMemoryAttrs(ze_context_handle_t hContext, void *ptr,
+               ze_device_handle_t *hDevice,
+               ZeStruct<ze_memory_allocation_properties_t> *properties);
 
 struct ze_region_params {
   const ze_copy_region_t dstRegion;

source/adapters/level_zero/memory.cpp

@@ -1746,28 +1746,19 @@ ur_result_t urMemBufferCreateWithNativeHandle(
   std::shared_lock<ur_shared_mutex> Lock(Context->Mutex);
 
   // Get base of the allocation
-  void *Base = nullptr;
-  size_t Size = 0;
   void *Ptr = ur_cast<void *>(NativeMem);
+
+  void *Base;
+  size_t Size;
   ZE2UR_CALL(zeMemGetAddressRange, (Context->ZeContext, Ptr, &Base, &Size));
-  UR_ASSERT(Ptr == Base, UR_RESULT_ERROR_INVALID_VALUE);
 
-  ZeStruct<ze_memory_allocation_properties_t> ZeMemProps;
-  ze_device_handle_t ZeDevice = nullptr;
-  ZE2UR_CALL(zeMemGetAllocProperties,
-             (Context->ZeContext, Ptr, &ZeMemProps, &ZeDevice));
+  UR_ASSERT(Ptr == Base, UR_RESULT_ERROR_INVALID_VALUE);
 
-  // Check type of the allocation
-  switch (ZeMemProps.type) {
-  case ZE_MEMORY_TYPE_HOST:
-  case ZE_MEMORY_TYPE_SHARED:
-  case ZE_MEMORY_TYPE_DEVICE:
-    break;
-  case ZE_MEMORY_TYPE_UNKNOWN:
-    // Memory allocation is unrelated to the context
-    return UR_RESULT_ERROR_INVALID_CONTEXT;
-  default:
-    die("Unexpected memory type");
+  ze_device_handle_t ZeDevice;
+  ZeStruct<ze_memory_allocation_properties_t> MemoryAttrs;
+  UR_CALL(getMemoryAttrs(Context->ZeContext, Ptr, &ZeDevice, &MemoryAttrs));
+  if (MemoryAttrs.type == ZE_MEMORY_TYPE_UNKNOWN) {
+    return UR_RESULT_ERROR_INVALID_VALUE;
   }
 
   ur_device_handle_t Device{};

source/adapters/level_zero/v2/api.cpp

@@ -18,20 +18,6 @@
 std::mutex ZeCall::GlobalLock;
 
 namespace ur::level_zero {
-ur_result_t urContextGetNativeHandle(ur_context_handle_t hContext,
-                                     ur_native_handle_t *phNativeContext) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
-ur_result_t urContextCreateWithNativeHandle(
-    ur_native_handle_t hNativeContext, ur_adapter_handle_t hAdapter,
-    uint32_t numDevices, const ur_device_handle_t *phDevices,
-    const ur_context_native_properties_t *pProperties,
-    ur_context_handle_t *phContext) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
 
 ur_result_t
 urContextSetExtendedDeleter(ur_context_handle_t hContext,
@@ -49,13 +35,6 @@ ur_result_t urMemImageCreate(ur_context_handle_t hContext, ur_mem_flags_t flags,
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
-ur_result_t urMemGetNativeHandle(ur_mem_handle_t hMem,
-                                 ur_device_handle_t hDevice,
-                                 ur_native_handle_t *phNativeMem) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
 ur_result_t urMemImageCreateWithNativeHandle(
     ur_native_handle_t hNativeMem, ur_context_handle_t hContext,
     const ur_image_format_t *pImageFormat, const ur_image_desc_t *pImageDesc,
@@ -192,22 +171,6 @@ ur_result_t urKernelSetSpecializationConstants(
   return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
 }
 
-ur_result_t urKernelGetNativeHandle(ur_kernel_handle_t hKernel,
-                                    ur_native_handle_t *phNativeKernel) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
-ur_result_t
-urKernelCreateWithNativeHandle(ur_native_handle_t hNativeKernel,
-                               ur_context_handle_t hContext,
-                               ur_program_handle_t hProgram,
-                               const ur_kernel_native_properties_t *pProperties,
-                               ur_kernel_handle_t *phKernel) {
-  logger::error("{} function not implemented!", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-}
-
 ur_result_t urKernelGetSuggestedLocalWorkSize(ur_kernel_handle_t hKernel,
                                               ur_queue_handle_t hQueue,
                                               uint32_t numWorkDim,

source/adapters/level_zero/v2/command_list_cache.cpp

@@ -76,8 +76,7 @@ command_list_cache_t::createCommandList(const command_list_descriptor_t &desc) {
   }
 }
 
-raii::cache_borrowed_command_list_t
-command_list_cache_t::getImmediateCommandList(
+raii::command_list_unique_handle command_list_cache_t::getImmediateCommandList(
     ze_device_handle_t ZeDevice, bool IsInOrder, uint32_t Ordinal,
     ze_command_queue_mode_t Mode, ze_command_queue_priority_t Priority,
     std::optional<uint32_t> Index) {
@@ -91,14 +90,14 @@ command_list_cache_t::getImmediateCommandList(
   Desc.Priority = Priority;
   Desc.Index = Index;
 
-  auto CommandList = getCommandList(Desc).release();
-  return raii::cache_borrowed_command_list_t(
+  auto [CommandList, _] = getCommandList(Desc).release();
+  return raii::command_list_unique_handle(
       CommandList, [Cache = this, Desc](ze_command_list_handle_t CmdList) {
         Cache->addCommandList(Desc, raii::ze_command_list_handle_t(CmdList));
       });
 }
 
-raii::cache_borrowed_command_list_t
+raii::command_list_unique_handle
 command_list_cache_t::getRegularCommandList(ze_device_handle_t ZeDevice,
                                             bool IsInOrder, uint32_t Ordinal) {
   TRACK_SCOPE_LATENCY("command_list_cache_t::getRegularCommandList");
@@ -108,9 +107,9 @@ command_list_cache_t::getRegularCommandList(ze_device_handle_t ZeDevice,
   Desc.IsInOrder = IsInOrder;
   Desc.Ordinal = Ordinal;
 
-  auto CommandList = getCommandList(Desc).release();
+  auto [CommandList, _] = getCommandList(Desc).release();
 
-  return raii::cache_borrowed_command_list_t(
+  return raii::command_list_unique_handle(
       CommandList, [Cache = this, Desc](ze_command_list_handle_t CmdList) {
         Cache->addCommandList(Desc, raii::ze_command_list_handle_t(CmdList));
       });

source/adapters/level_zero/v2/command_list_cache.hpp

@@ -21,7 +21,7 @@
 
 namespace v2 {
 namespace raii {
-using cache_borrowed_command_list_t =
+using command_list_unique_handle =
     std::unique_ptr<::_ze_command_list_handle_t,
                     std::function<void(::ze_command_list_handle_t)>>;
 } // namespace raii
@@ -54,12 +54,12 @@ struct command_list_descriptor_hash_t {
 struct command_list_cache_t {
   command_list_cache_t(ze_context_handle_t ZeContext);
 
-  raii::cache_borrowed_command_list_t
+  raii::command_list_unique_handle
   getImmediateCommandList(ze_device_handle_t ZeDevice, bool IsInOrder,
                           uint32_t Ordinal, ze_command_queue_mode_t Mode,
                           ze_command_queue_priority_t Priority,
                           std::optional<uint32_t> Index = std::nullopt);
-  raii::cache_borrowed_command_list_t
+  raii::command_list_unique_handle
   getRegularCommandList(ze_device_handle_t ZeDevice, bool IsInOrder,
                         uint32_t Ordinal);
 

source/adapters/level_zero/v2/common.hpp

@@ -64,18 +64,20 @@ struct ze_handle_wrapper {
       return;
     }
 
-    auto zeResult = ZE_CALL_NOCHECK(destroy, (handle));
-    // Gracefully handle the case that L0 was already unloaded.
-    if (zeResult && zeResult != ZE_RESULT_ERROR_UNINITIALIZED)
-      throw ze2urResult(zeResult);
+    if (ownZeHandle) {
+      auto zeResult = ZE_CALL_NOCHECK(destroy, (handle));
+      // Gracefully handle the case that L0 was already unloaded.
+      if (zeResult && zeResult != ZE_RESULT_ERROR_UNINITIALIZED)
+        throw ze2urResult(zeResult);
+    }
 
     handle = nullptr;
   }
 
-  ZeHandleT release() {
+  std::pair<ZeHandleT, bool> release() {
     auto handle = this->handle;
     this->handle = nullptr;
-    return handle;
+    return {handle, ownZeHandle};
   }
 
   ZeHandleT get() const { return handle; }

source/adapters/level_zero/v2/context.cpp

@@ -120,6 +120,27 @@ ur_result_t urContextCreate(uint32_t deviceCount,
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t urContextGetNativeHandle(ur_context_handle_t hContext,
+                                     ur_native_handle_t *phNativeContext) {
+  *phNativeContext =
+      reinterpret_cast<ur_native_handle_t>(hContext->getZeHandle());
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t urContextCreateWithNativeHandle(
+    ur_native_handle_t hNativeContext, ur_adapter_handle_t, uint32_t numDevices,
+    const ur_device_handle_t *phDevices,
+    const ur_context_native_properties_t *pProperties,
+    ur_context_handle_t *phContext) {
+  auto zeContext = reinterpret_cast<ze_context_handle_t>(hNativeContext);
+
+  auto ownZeHandle = pProperties ? pProperties->isNativeHandleOwned : false;
+
+  *phContext =
+      new ur_context_handle_t_(zeContext, numDevices, phDevices, ownZeHandle);
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t urContextRetain(ur_context_handle_t hContext) {
   return hContext->retain();
 }
@@ -134,7 +155,8 @@ ur_result_t urContextGetInfo(ur_context_handle_t hContext,
                              void *pContextInfo,
 
                              size_t *pPropSizeRet) {
-  std::shared_lock<ur_shared_mutex> Lock(hContext->Mutex);
+  // No locking needed here, we only read const members
+
   UrReturnHelper ReturnValue(propSize, pContextInfo, pPropSizeRet);
   switch (
       (uint32_t)contextInfoType) { // cast to avoid warnings on EXT enum values

source/adapters/level_zero/v2/event.cpp

@@ -93,8 +93,6 @@ static uint64_t adjustEndEventTimestamp(uint64_t adjustedStartTimestamp,
 }
 
 uint64_t ur_event_handle_t_::getEventEndTimestamp() {
-  std::scoped_lock<ur_shared_mutex> lock(this->Mutex);
-
   // If adjustedEventEndTimestamp on the event is non-zero it means it has
   // collected the result of the queue already. In that case it has been
   // adjusted and is ready for immediate return.
@@ -120,8 +118,6 @@ void ur_event_handle_t_::recordStartTimestamp() {
   UR_CALL_THROWS(ur::level_zero::urDeviceGetGlobalTimestamps(
       getDevice(), &deviceStartTimestamp, nullptr));
 
-  std::scoped_lock<ur_shared_mutex> lock(this->Mutex);
-
   adjustedEventStartTimestamp = deviceStartTimestamp;
 }
 
@@ -183,6 +179,8 @@ ur_result_t urEventGetProfilingInfo(
     size_t *pPropValueSizeRet ///< [out][optional] pointer to the actual size in
                               ///< bytes returned in propValue
 ) {
+  std::scoped_lock<ur_shared_mutex> lock(hEvent->Mutex);
+
   // The event must either have profiling enabled or be recording timestamps.
   bool isTimestampedEvent = hEvent->isTimestamped();
   if (!hEvent->isProfilingEnabled() && !isTimestampedEvent) {

source/adapters/level_zero/v2/event_provider_counter.cpp

@@ -54,7 +54,7 @@ raii::cache_borrowed_event provider_counter::allocate() {
   freelist.pop_back();
 
   return raii::cache_borrowed_event(
-      event.release(),
+      event.release().first,
       [this](ze_event_handle_t handle) { freelist.push_back(handle); });
 }
 

source/adapters/level_zero/v2/event_provider_normal.cpp

@@ -68,7 +68,7 @@ raii::cache_borrowed_event provider_pool::allocate() {
   auto e = std::move(freelist.back());
   freelist.pop_back();
   return raii::cache_borrowed_event(
-      e.release(),
+      e.release().first,
       [this](ze_event_handle_t handle) { freelist.push_back(handle); });
 }