Merge pull request #2318 from pbalcer/fix-native-kernel-handles

[L0 v2] fix native kernel handle interop

Author: pbalcer

source/adapters/level_zero/v2/kernel.cpp

@@ -70,8 +70,10 @@ ur_kernel_handle_t_::ur_kernel_handle_t_(ur_program_handle_t hProgram,
 
 ur_kernel_handle_t_::ur_kernel_handle_t_(
     ur_native_handle_t hNativeKernel, ur_program_handle_t hProgram,
+    ur_context_handle_t context,
     const ur_kernel_native_properties_t *pProperties)
-    : hProgram(hProgram), deviceKernels(1) {
+    : hProgram(hProgram),
+      deviceKernels(context ? context->getPlatform()->getNumDevices() : 0) {
   ur::level_zero::urProgramRetain(hProgram);
 
   auto ownZeHandle = pProperties ? pProperties->isNativeHandleOwned : false;
@@ -82,7 +84,12 @@ ur_kernel_handle_t_::ur_kernel_handle_t_(
     throw UR_RESULT_ERROR_INVALID_KERNEL;
   }
 
-  deviceKernels.back().emplace(nullptr, zeKernel, ownZeHandle);
+  for (auto &Dev : context->getDevices()) {
+    deviceKernels[*Dev->Id].emplace(Dev, zeKernel, ownZeHandle);
+
+    // owned only by the first entry
+    ownZeHandle = false;
+  }
   completeInitialization();
 }
 
@@ -128,20 +135,6 @@ size_t ur_kernel_handle_t_::deviceIndex(ur_device_handle_t hDevice) const {
     hDevice = hDevice->RootDevice;
   }
 
-  // supports kernels created from native handle
-  if (deviceKernels.size() == 1) {
-    assert(deviceKernels[0].has_value());
-    assert(deviceKernels[0].value().hKernel.get());
-
-    auto &kernel = deviceKernels[0].value();
-
-    if (kernel.hDevice != hDevice) {
-      throw UR_RESULT_ERROR_INVALID_DEVICE;
-    }
-
-    return 0;
-  }
-
   if (!deviceKernels[hDevice->Id.value()].has_value()) {
     throw UR_RESULT_ERROR_INVALID_DEVICE;
   }
@@ -341,8 +334,12 @@ urKernelCreateWithNativeHandle(ur_native_handle_t hNativeKernel,
                                ur_program_handle_t hProgram,
                                const ur_kernel_native_properties_t *pProperties,
                                ur_kernel_handle_t *phKernel) {
-  std::ignore = hContext;
-  *phKernel = new ur_kernel_handle_t_(hNativeKernel, hProgram, pProperties);
+  if (!hProgram) {
+    return UR_RESULT_ERROR_INVALID_NULL_HANDLE;
+  }
+
+  *phKernel =
+      new ur_kernel_handle_t_(hNativeKernel, hProgram, hContext, pProperties);
   return UR_RESULT_SUCCESS;
 }
 

source/adapters/level_zero/v2/kernel.hpp

@@ -39,7 +39,7 @@ struct ur_kernel_handle_t_ : _ur_object {
 
   // From native handle
   ur_kernel_handle_t_(ur_native_handle_t hNativeKernel,
-                      ur_program_handle_t hProgram,
+                      ur_program_handle_t hProgram, ur_context_handle_t context,
                       const ur_kernel_native_properties_t *pProperties);
 
   // Get L0 kernel handle for a given device

test/adapters/level_zero/urKernelCreateWithNativeHandle.cpp

@@ -8,7 +8,7 @@
 #include "ze_api.h"
 #include <uur/fixtures.h>
 
-using urLevelZeroKernelNativeHandleTest = uur::urContextTest;
+using urLevelZeroKernelNativeHandleTest = uur::urQueueTest;
 UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urLevelZeroKernelNativeHandleTest);
 
 TEST_P(urLevelZeroKernelNativeHandleTest, OwnedHandleRelease) {
@@ -60,6 +60,13 @@ TEST_P(urLevelZeroKernelNativeHandleTest, OwnedHandleRelease) {
     ASSERT_SUCCESS(urKernelCreateWithNativeHandle(
         (ur_native_handle_t)native_kernel, context, program, &kprops, &kernel));
 
+    size_t global_offset = 0;
+    size_t local_size = 1;
+    size_t global_size = 1;
+    ASSERT_SUCCESS(urEnqueueKernelLaunch(queue, kernel, 1, &global_offset,
+                                         &local_size, &global_size, 0, nullptr,
+                                         nullptr));
+
     ASSERT_SUCCESS(urKernelRelease(kernel));
     ASSERT_SUCCESS(urProgramRelease(program));
 }