[SYCL] Add ownership control for LeveL-Zero kernel_bundle interop.

Signed-off-by: Sergey V Maslov <[email protected]>

Author: smaslov-intel

sycl/include/CL/sycl/backend.hpp

@@ -108,9 +108,13 @@ __SYCL_EXPORT event make_event(pi_native_handle NativeHandle,
                                backend Backend);
 __SYCL_EXPORT kernel make_kernel(pi_native_handle NativeHandle,
                                  const context &TargetContext, backend Backend);
+// TODO: Unused. Remove when allowed.
 __SYCL_EXPORT std::shared_ptr<detail::kernel_bundle_impl>
 make_kernel_bundle(pi_native_handle NativeHandle, const context &TargetContext,
                    bundle_state State, backend Backend);
+__SYCL_EXPORT std::shared_ptr<detail::kernel_bundle_impl>
+make_kernel_bundle(pi_native_handle NativeHandle, const context &TargetContext,
+                   bool KeepOwnership, bundle_state State, backend Backend);
 } // namespace detail
 
 template <backend Backend>
@@ -221,7 +225,7 @@ make_kernel_bundle(const typename backend_traits<Backend>::template input_type<
   std::shared_ptr<detail::kernel_bundle_impl> KBImpl =
       detail::make_kernel_bundle(
           detail::pi::cast<pi_native_handle>(BackendObject), TargetContext,
-          State, Backend);
+          false, State, Backend);
   return detail::createSyclObjFromImpl<kernel_bundle<State>>(KBImpl);
 }
 } // namespace sycl

sycl/include/CL/sycl/detail/pi.h

@@ -40,6 +40,7 @@
 // changes the API version from 3.5 to 4.6.
 // 5.7 Added new context and ownership arguments to
 //   piextEventCreateWithNativeHandle
+// 6.8 Added new ownership argument to piextProgramCreateWithNativeHandle.
 //
 #include "CL/cl.h"
 #define _PI_H_VERSION_MAJOR 5
@@ -1219,9 +1220,12 @@ piextProgramGetNativeHandle(pi_program program, pi_native_handle *nativeHandle);
 ///
 /// \param nativeHandle is the native handle to create PI program from.
 /// \param context is the PI context of the program.
+/// \param ownNativeHandle tells if SYCL RT should assume the ownership of
+///        the native handle, if it can.
 /// \param program is the PI program created from the native handle.
 __SYCL_EXPORT pi_result piextProgramCreateWithNativeHandle(
-    pi_native_handle nativeHandle, pi_context context, pi_program *program);
+    pi_native_handle nativeHandle, pi_context context, bool ownNativeHandle,
+    pi_program *program);
 
 //
 // Kernel

sycl/include/sycl/ext/oneapi/backend/level_zero.hpp

@@ -101,7 +101,11 @@ template <> struct BackendInput<backend::level_zero, event> {
 
 template <bundle_state State>
 struct BackendInput<backend::level_zero, kernel_bundle<State>> {
-  using type = ze_module_handle_t;
+  using type = struct {
+    ze_module_handle_t NativeHandle;
+    ext::oneapi::level_zero::ownership Ownership{
+        ext::oneapi::level_zero::ownership::transfer};
+  };
 };
 
 template <bundle_state State>
@@ -249,6 +253,24 @@ event make_event<backend::level_zero>(
       BackendObject.Ownership == ext::oneapi::level_zero::ownership::keep);
 }
 
+// Specialization of sycl::make_kernel_bundle for Level-Zero backend.
+template <>
+kernel_bundle<bundle_state::executable>
+make_kernel_bundle<backend::ext_oneapi_level_zero, bundle_state::executable>(
+    const backend_input_t<backend::ext_oneapi_level_zero,
+                          kernel_bundle<bundle_state::executable>>
+        &BackendObject,
+    const context &TargetContext) {
+  std::shared_ptr<detail::kernel_bundle_impl> KBImpl =
+      detail::make_kernel_bundle(
+          detail::pi::cast<pi_native_handle>(BackendObject.NativeHandle),
+          TargetContext,
+          BackendObject.Ownership == ext::oneapi::level_zero::ownership::keep,
+          bundle_state::executable, backend::ext_oneapi_level_zero);
+  return detail::createSyclObjFromImpl<kernel_bundle<bundle_state::executable>>(
+      KBImpl);
+}
+
 // TODO: remove this specialization when generic is changed to call
 // .GetNative() instead of .get_native() member of kernel_bundle.
 template <>

sycl/plugins/cuda/pi_cuda.cpp

@@ -3173,7 +3173,7 @@ pi_result cuda_piextProgramGetNativeHandle(pi_program program,
 ///
 /// \return TBD
 pi_result cuda_piextProgramCreateWithNativeHandle(pi_native_handle, pi_context,
-                                                  pi_program *) {
+                                                  bool, pi_program *) {
   cl::sycl::detail::pi::die(
       "Creation of PI program from native handle not implemented");
   return {};

sycl/plugins/esimd_cpu/pi_esimd_cpu.cpp

@@ -1019,19 +1019,15 @@ pi_result piProgramGetBuildInfo(pi_program, pi_device, cl_program_build_info,
   DIE_NO_IMPLEMENTATION;
 }
 
-pi_result piProgramRetain(pi_program) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piProgramRetain(pi_program) { DIE_NO_IMPLEMENTATION; }
 
-pi_result piProgramRelease(pi_program) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piProgramRelease(pi_program) { DIE_NO_IMPLEMENTATION; }
 
 pi_result piextProgramGetNativeHandle(pi_program, pi_native_handle *) {
   DIE_NO_IMPLEMENTATION;
 }
 
-pi_result piextProgramCreateWithNativeHandle(pi_native_handle, pi_context,
+pi_result piextProgramCreateWithNativeHandle(pi_native_handle, pi_context, bool,
                                              pi_program *) {
   DIE_NO_IMPLEMENTATION;
 }
@@ -1068,17 +1064,11 @@ pi_result piKernelGetSubGroupInfo(pi_kernel, pi_device,
   DIE_NO_IMPLEMENTATION;
 }
 
-pi_result piKernelRetain(pi_kernel) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piKernelRetain(pi_kernel) { DIE_NO_IMPLEMENTATION; }
 
-pi_result piKernelRelease(pi_kernel) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piKernelRelease(pi_kernel) { DIE_NO_IMPLEMENTATION; }
 
-pi_result piEventCreate(pi_context, pi_event *) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piEventCreate(pi_context, pi_event *) { DIE_NO_IMPLEMENTATION; }
 
 pi_result piEventGetInfo(pi_event, pi_event_info, size_t, void *, size_t *) {
   DIE_NO_IMPLEMENTATION;
@@ -1117,9 +1107,7 @@ pi_result piEventSetCallback(pi_event, pi_int32,
   DIE_NO_IMPLEMENTATION;
 }
 
-pi_result piEventSetStatus(pi_event, pi_int32) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piEventSetStatus(pi_event, pi_int32) { DIE_NO_IMPLEMENTATION; }
 
 pi_result piEventRetain(pi_event Event) {
   if (Event == nullptr) {
@@ -1170,13 +1158,9 @@ pi_result piSamplerGetInfo(pi_sampler, pi_sampler_info, size_t, void *,
   DIE_NO_IMPLEMENTATION;
 }
 
-pi_result piSamplerRetain(pi_sampler) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piSamplerRetain(pi_sampler) { DIE_NO_IMPLEMENTATION; }
 
-pi_result piSamplerRelease(pi_sampler) {
-  DIE_NO_IMPLEMENTATION;
-}
+pi_result piSamplerRelease(pi_sampler) { DIE_NO_IMPLEMENTATION; }
 
 pi_result piEnqueueEventsWait(pi_queue, pi_uint32, const pi_event *,
                               pi_event *) {

sycl/plugins/hip/pi_hip.cpp

@@ -2949,11 +2949,14 @@ pi_result hip_piextProgramGetNativeHandle(pi_program program,
 ///
 /// \param[in] nativeHandle The native handle to create PI program object from.
 /// \param[in] context The PI context of the program.
+/// \param[in] ownNativeHandle tells if should assume the ownership of
+///            the native handle.
 /// \param[out] program Set to the PI program object created from native handle.
 ///
 /// \return TBD
 pi_result hip_piextProgramCreateWithNativeHandle(pi_native_handle nativeHandle,
                                                  pi_context context,
+                                                 bool ownNativeHandle,
                                                  pi_program *program) {
   cl::sycl::detail::pi::die(
       "Creation of PI program from native handle not implemented");

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -3657,8 +3657,9 @@ pi_result piProgramLink(pi_context Context, pi_uint32 NumDevices,
           if (res != PI_SUCCESS) {
             return res;
           }
-          Input = new _pi_program(Input->Context, ZeModule, _pi_program::Object,
-                                  Input->HasImports);
+          Input =
+              new _pi_program(Input->Context, ZeModule, true /*own ZeModule*/,
+                              _pi_program::Object, Input->HasImports);
           Input->HasImportsAndIsLinked = true;
         }
       } else {
@@ -3913,6 +3914,7 @@ pi_result piextProgramGetNativeHandle(pi_program Program,
 
 pi_result piextProgramCreateWithNativeHandle(pi_native_handle NativeHandle,
                                              pi_context Context,
+                                             bool ownNativeHandle,
                                              pi_program *Program) {
   PI_ASSERT(Program, PI_INVALID_PROGRAM);
   PI_ASSERT(NativeHandle, PI_INVALID_VALUE);
@@ -3925,7 +3927,8 @@ pi_result piextProgramCreateWithNativeHandle(pi_native_handle NativeHandle,
   // executable (state Object).
 
   try {
-    *Program = new _pi_program(Context, ZeModule, _pi_program::Exe);
+    *Program =
+        new _pi_program(Context, ZeModule, ownNativeHandle, _pi_program::Exe);
   } catch (const std::bad_alloc &) {
     return PI_OUT_OF_HOST_MEMORY;
   } catch (...) {
@@ -3942,7 +3945,7 @@ _pi_program::~_pi_program() {
     ZE_CALL_NOCHECK(zeModuleBuildLogDestroy, (ZeBuildLog));
   }
 
-  if (ZeModule) {
+  if (ZeModule && OwnZeModule) {
     ZE_CALL_NOCHECK(zeModuleDestroy, (ZeModule));
   }
 }

sycl/plugins/level_zero/pi_level_zero.hpp

@@ -1068,23 +1068,24 @@ struct _pi_program : _pi_object {
   // Construct a program in IL or Native state.
   _pi_program(pi_context Context, const void *Input, size_t Length, state St)
       : State(St), Context(Context), Code(new uint8_t[Length]),
-        CodeLength(Length), ZeModule(nullptr), HasImports(false),
-        HasImportsAndIsLinked(false), ZeBuildLog(nullptr) {
+        CodeLength(Length), ZeModule(nullptr), OwnZeModule{true},
+        HasImports(false), HasImportsAndIsLinked(false), ZeBuildLog(nullptr) {
 
     std::memcpy(Code.get(), Input, Length);
   }
 
   // Construct a program in either Object or Exe state.
-  _pi_program(pi_context Context, ze_module_handle_t ZeModule, state St,
-              bool HasImports = false)
-      : State(St), Context(Context), ZeModule(ZeModule), HasImports(HasImports),
+  _pi_program(pi_context Context, ze_module_handle_t ZeModule, bool OwnZeModule,
+              state St, bool HasImports = false)
+      : State(St), Context(Context),
+        ZeModule(ZeModule), OwnZeModule{OwnZeModule}, HasImports(HasImports),
         HasImportsAndIsLinked(false), ZeBuildLog(nullptr) {}
 
   // Construct a program in LinkedExe state.
   _pi_program(pi_context Context, std::vector<LinkedReleaser> &&Inputs,
               ze_module_build_log_handle_t ZeLog)
       : State(LinkedExe), Context(Context), ZeModule(nullptr),
-        HasImports(false), HasImportsAndIsLinked(false),
+        OwnZeModule(true), HasImports(false), HasImportsAndIsLinked(false),
         LinkedPrograms(std::move(Inputs)), ZeBuildLog(ZeLog) {}
 
   ~_pi_program();
@@ -1103,7 +1104,13 @@ struct _pi_program : _pi_object {
 
   // Used for programs in Object or Exe state.
   ze_module_handle_t ZeModule; // Level Zero module handle.
-  bool HasImports;             // Tells if module imports any symbols.
+
+  // Indicates if we own the ZeModule or it came from interop that
+  // asked to not transfer the ownership to SYCL RT.
+  bool OwnZeModule;
+
+  // Tells if module imports any symbols.
+  bool HasImports;
 
   // Used for programs in Object state.  Tells if this module imports any
   // symbols AND it is linked into some other program that has state LinkedExe.

sycl/plugins/opencl/pi_opencl.cpp

@@ -444,7 +444,7 @@ pi_result piProgramCreate(pi_context context, const void *il, size_t length,
 }
 
 pi_result piextProgramCreateWithNativeHandle(pi_native_handle nativeHandle,
-                                             pi_context,
+                                             pi_context, bool,
                                              pi_program *piProgram) {
   assert(piProgram != nullptr);
   *piProgram = reinterpret_cast<pi_program>(nativeHandle);

sycl/source/backend.cpp

@@ -118,13 +118,13 @@ __SYCL_EXPORT event make_event(pi_native_handle NativeHandle,
 
 std::shared_ptr<detail::kernel_bundle_impl>
 make_kernel_bundle(pi_native_handle NativeHandle, const context &TargetContext,
-                   bundle_state State, backend Backend) {
+                   bool KeepOwnership, bundle_state State, backend Backend) {
   const auto &Plugin = getPlugin(Backend);
   const auto &ContextImpl = getSyclObjImpl(TargetContext);
 
   pi::PiProgram PiProgram = nullptr;
   Plugin.call<PiApiKind::piextProgramCreateWithNativeHandle>(
-      NativeHandle, ContextImpl->getHandleRef(), &PiProgram);
+      NativeHandle, ContextImpl->getHandleRef(), KeepOwnership, &PiProgram);
 
   std::vector<pi::PiDevice> ProgramDevices;
   size_t NumDevices = 0;
@@ -194,6 +194,14 @@ make_kernel_bundle(pi_native_handle NativeHandle, const context &TargetContext,
 
   return std::make_shared<kernel_bundle_impl>(TargetContext, Devices, DevImg);
 }
+
+// TODO: Unused. Remove when allowed.
+std::shared_ptr<detail::kernel_bundle_impl>
+make_kernel_bundle(pi_native_handle NativeHandle, const context &TargetContext,
+                   bundle_state State, backend Backend) {
+  return make_kernel_bundle(NativeHandle, TargetContext, false, State, Backend);
+}
+
 kernel make_kernel(pi_native_handle NativeHandle, const context &TargetContext,
                    backend Backend) {
   const auto &Plugin = getPlugin(Backend);

sycl/source/detail/program_impl.cpp

@@ -129,7 +129,7 @@ program_impl::program_impl(ContextImplPtr Context,
            "No InteropProgram/PiProgram defined with piextProgramFromNative");
     // Translate the raw program handle into PI program.
     Plugin.call<PiApiKind::piextProgramCreateWithNativeHandle>(
-        InteropProgram, MContext->getHandleRef(), &MProgram);
+        InteropProgram, MContext->getHandleRef(), false, &MProgram);
   } else
     Plugin.call<PiApiKind::piProgramRetain>(Program);
 

sycl/test/basic_tests/interop-level-zero-2020.cpp

@@ -137,8 +137,9 @@ int main() {
       {ZeEvent, ext::oneapi::level_zero::ownership::keep}, Context);
   kernel_bundle<bundle_state::executable> InteropKernelBundle =
       make_kernel_bundle<backend::ext_oneapi_level_zero,
-                         bundle_state::executable>(ZeKernelBundle.front(),
-                                                   Context);
+                         bundle_state::executable>(
+          {ZeKernelBundle.front(), ext::oneapi::level_zero::ownership::keep},
+          Context);
 
   // Check deprecated
   // expected-warning@+1 {{'make<sycl::platform, nullptr>' is deprecated: Use SYCL 2020 sycl::make_platform free function}}