[L0 v2] implement native APIs for kernel, context and memory

Author: igchor

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

@@ -91,7 +91,7 @@ command_list_cache_t::getImmediateCommandList(
   Desc.Priority = Priority;
   Desc.Index = Index;
 
-  auto CommandList = getCommandList(Desc).release();
+  auto [CommandList, _] = getCommandList(Desc).release();
   return raii::cache_borrowed_command_list_t(
       CommandList, [Cache = this, Desc](ze_command_list_handle_t CmdList) {
         Cache->addCommandList(Desc, raii::ze_command_list_handle_t(CmdList));
@@ -108,7 +108,7 @@ 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(
       CommandList, [Cache = this, Desc](ze_command_list_handle_t CmdList) {

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();
 }

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); });
 }