Merge pull request #1778 from nrspruit/fix_multi_driver_exp_handle

[L0] Fix Handle used in calls to L0 Driver zex apis given multi drivers

Author: callumfare

source/adapters/level_zero/adapter.cpp

@@ -49,6 +49,9 @@ ur_result_t initPlatforms(PlatformVec &platforms) noexcept try {
   for (uint32_t I = 0; I < ZeDriverCount; ++I) {
     auto platform = std::make_unique<ur_platform_handle_t_>(ZeDrivers[I]);
     UR_CALL(platform->initialize());
+    ZE2UR_CALL(zelLoaderTranslateHandle,
+               (ZEL_HANDLE_DRIVER, platform->ZeDriver,
+                (void **)&platform->ZeDriverHandleExpTranslated));
 
     // Save a copy in the cache for future uses.
     platforms.push_back(std::move(platform));

source/adapters/level_zero/adapter.hpp

@@ -11,6 +11,7 @@
 
 #include "logger/ur_logger.hpp"
 #include <atomic>
+#include <loader/ze_loader.h>
 #include <mutex>
 #include <optional>
 #include <ur/ur.hpp>

source/adapters/level_zero/command_buffer.cpp

@@ -56,12 +56,14 @@ bool PreferCopyEngineForFill = [] {
 ur_exp_command_buffer_handle_t_::ur_exp_command_buffer_handle_t_(
     ur_context_handle_t Context, ur_device_handle_t Device,
     ze_command_list_handle_t CommandList,
+    ze_command_list_handle_t CommandListTranslated,
     ze_command_list_handle_t CommandListResetEvents,
     ze_command_list_handle_t CopyCommandList,
     ZeStruct<ze_command_list_desc_t> ZeDesc,
     ZeStruct<ze_command_list_desc_t> ZeCopyDesc,
     const ur_exp_command_buffer_desc_t *Desc, const bool IsInOrderCmdList)
     : Context(Context), Device(Device), ZeComputeCommandList(CommandList),
+      ZeComputeCommandListTranslated(CommandListTranslated),
       ZeCommandListResetEvents(CommandListResetEvents),
       ZeCommandListDesc(ZeDesc), ZeCopyCommandList(CopyCommandList),
       ZeCopyCommandListDesc(ZeCopyDesc), ZeFencesMap(), ZeActiveFence(nullptr),
@@ -605,11 +607,16 @@ urCommandBufferCreateExp(ur_context_handle_t Context, ur_device_handle_t Device,
                 &ZeCopyCommandList));
   }
 
+  ze_command_list_handle_t ZeComputeCommandListTranslated = nullptr;
+  ZE2UR_CALL(zelLoaderTranslateHandle,
+             (ZEL_HANDLE_COMMAND_LIST, ZeComputeCommandList,
+              (void **)&ZeComputeCommandListTranslated));
+
   try {
     *CommandBuffer = new ur_exp_command_buffer_handle_t_(
-        Context, Device, ZeComputeCommandList, ZeCommandListResetEvents,
-        ZeCopyCommandList, ZeCommandListDesc, ZeCopyCommandListDesc,
-        CommandBufferDesc, IsInOrder);
+        Context, Device, ZeComputeCommandList, ZeComputeCommandListTranslated,
+        ZeCommandListResetEvents, ZeCopyCommandList, ZeCommandListDesc,
+        ZeCopyCommandListDesc, CommandBufferDesc, IsInOrder);
   } catch (const std::bad_alloc &) {
     return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
   } catch (...) {
@@ -791,8 +798,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferAppendKernelLaunchExp(
     UR_ASSERT(Plt->ZeMutableCmdListExt.Supported,
               UR_RESULT_ERROR_UNSUPPORTED_FEATURE);
     ZE2UR_CALL(Plt->ZeMutableCmdListExt.zexCommandListGetNextCommandIdExp,
-               (CommandBuffer->ZeComputeCommandList, &ZeMutableCommandDesc,
-                &CommandId));
+               (CommandBuffer->ZeComputeCommandListTranslated,
+                &ZeMutableCommandDesc, &CommandId));
     DEBUG_LOG(CommandId);
   }
   try {
@@ -1619,8 +1626,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferUpdateKernelLaunchExp(
   auto Plt = CommandBuffer->Context->getPlatform();
   UR_ASSERT(Plt->ZeMutableCmdListExt.Supported,
             UR_RESULT_ERROR_UNSUPPORTED_FEATURE);
-  ZE2UR_CALL(Plt->ZeMutableCmdListExt.zexCommandListUpdateMutableCommandsExp,
-             (CommandBuffer->ZeComputeCommandList, &MutableCommandDesc));
+  ZE2UR_CALL(
+      Plt->ZeMutableCmdListExt.zexCommandListUpdateMutableCommandsExp,
+      (CommandBuffer->ZeComputeCommandListTranslated, &MutableCommandDesc));
   ZE2UR_CALL(zeCommandListClose, (CommandBuffer->ZeComputeCommandList));
 
   return UR_RESULT_SUCCESS;

source/adapters/level_zero/command_buffer.hpp

@@ -28,6 +28,7 @@ struct ur_exp_command_buffer_handle_t_ : public _ur_object {
   ur_exp_command_buffer_handle_t_(
       ur_context_handle_t Context, ur_device_handle_t Device,
       ze_command_list_handle_t CommandList,
+      ze_command_list_handle_t CommandListTranslated,
       ze_command_list_handle_t CommandListResetEvents,
       ze_command_list_handle_t CopyCommandList,
       ZeStruct<ze_command_list_desc_t> ZeDesc,
@@ -55,6 +56,9 @@ struct ur_exp_command_buffer_handle_t_ : public _ur_object {
   ur_device_handle_t Device;
   // Level Zero command list handle
   ze_command_list_handle_t ZeComputeCommandList;
+  // Given a multi driver scenario, the driver handle must be translated to the
+  // internal driver handle to allow calls to driver experimental apis.
+  ze_command_list_handle_t ZeComputeCommandListTranslated;
   // Level Zero command list handle
   ze_command_list_handle_t ZeCommandListResetEvents;
   // Level Zero command list descriptor

source/adapters/level_zero/image.cpp

@@ -488,7 +488,11 @@ ur_result_t bindlessImagesCreateImpl(ur_context_handle_t hContext,
     return UR_RESULT_ERROR_INVALID_OPERATION;
 
   uint64_t DeviceOffset{};
-  ZE2UR_CALL(zeImageGetDeviceOffsetExpFunctionPtr, (ZeImage, &DeviceOffset));
+  ze_image_handle_t ZeImageTranslated;
+  ZE2UR_CALL(zelLoaderTranslateHandle,
+             (ZEL_HANDLE_IMAGE, ZeImage, (void **)&ZeImageTranslated));
+  ZE2UR_CALL(zeImageGetDeviceOffsetExpFunctionPtr,
+             (ZeImageTranslated, &DeviceOffset));
   *phImage = reinterpret_cast<ur_exp_image_handle_t>(DeviceOffset);
 
   return UR_RESULT_SUCCESS;
@@ -652,8 +656,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMPitchedAllocExp(
 
   size_t Width = widthInBytes / elementSizeBytes;
   size_t RowPitch;
+  ze_device_handle_t ZeDeviceTranslated;
+  ZE2UR_CALL(zelLoaderTranslateHandle, (ZEL_HANDLE_DEVICE, hDevice->ZeDevice,
+                                        (void **)&ZeDeviceTranslated));
   ZE2UR_CALL(zeMemGetPitchFor2dImageFunctionPtr,
-             (hContext->ZeContext, hDevice->ZeDevice, Width, height,
+             (hContext->ZeContext, ZeDeviceTranslated, Width, height,
               elementSizeBytes, &RowPitch));
   *pResultPitch = RowPitch;
 

source/adapters/level_zero/memory.cpp

@@ -1668,7 +1668,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemBufferCreate(
     // If not shared of any type, we can import the ptr
     if (ZeMemoryAllocationProperties.type == ZE_MEMORY_TYPE_UNKNOWN) {
       // Promote the host ptr to USM host memory
-      ze_driver_handle_t driverHandle = Context->getPlatform()->ZeDriver;
+      ze_driver_handle_t driverHandle =
+          Context->getPlatform()->ZeDriverHandleExpTranslated;
       ZeUSMImport.doZeUSMImport(driverHandle, Host, Size);
       HostPtrImported = true;
     }
@@ -2252,7 +2253,8 @@ ur_result_t _ur_buffer::free() {
       UR_CALL(ZeMemFreeHelper(UrContext, ZeHandle));
       break;
     case allocation_t::unimport:
-      ZeUSMImport.doZeUSMRelease(UrContext->getPlatform()->ZeDriver, ZeHandle);
+      ZeUSMImport.doZeUSMRelease(
+          UrContext->getPlatform()->ZeDriverHandleExpTranslated, ZeHandle);
       break;
     default:
       die("_ur_buffer::free(): Unhandled release action");

source/adapters/level_zero/platform.hpp

@@ -24,6 +24,10 @@ struct ur_platform_handle_t_ : public _ur_platform {
   // a pretty good fit to keep here.
   ze_driver_handle_t ZeDriver;
 
+  // Given a multi driver scenario, the driver handle must be translated to the
+  // internal driver handle to allow calls to driver experimental apis.
+  ze_driver_handle_t ZeDriverHandleExpTranslated;
+
   // Cache versions info from zeDriverGetProperties.
   std::string ZeDriverVersion;
   std::string ZeDriverApiVersion;

source/adapters/level_zero/ur_level_zero.hpp

@@ -18,6 +18,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include <loader/ze_loader.h>
 #include <ur/ur.hpp>
 #include <ur_api.h>
 #include <ze_api.h>

source/adapters/level_zero/usm.cpp

@@ -1026,7 +1026,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMImportExp(ur_context_handle_t Context,
     // If not shared of any type, we can import the ptr
     if (ZeMemoryAllocationProperties.type == ZE_MEMORY_TYPE_UNKNOWN) {
       // Promote the host ptr to USM host memory
-      ze_driver_handle_t driverHandle = Context->getPlatform()->ZeDriver;
+      ze_driver_handle_t driverHandle =
+          Context->getPlatform()->ZeDriverHandleExpTranslated;
       ZeUSMImport.doZeUSMImport(driverHandle, HostPtr, Size);
     }
   }
@@ -1039,6 +1040,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urUSMReleaseExp(ur_context_handle_t Context,
 
   // Release the imported memory.
   if (ZeUSMImport.Supported && HostPtr != nullptr)
-    ZeUSMImport.doZeUSMRelease(Context->getPlatform()->ZeDriver, HostPtr);
+    ZeUSMImport.doZeUSMRelease(
+        Context->getPlatform()->ZeDriverHandleExpTranslated, HostPtr);
   return UR_RESULT_SUCCESS;
 }