[SYCL] enable_shared_from_this for kernel_bundle_impl

sycl/include/sycl/handler.hpp

@@ -891,7 +891,7 @@ class __SYCL_EXPORT handler {
     // If the kernel lambda is callable with a kernel_handler argument, manifest
     // the associated kernel handler.
     if constexpr (IsCallableWithKernelHandler) {
-      getOrInsertHandlerKernelBundle(/*Insert=*/true);
+      getOrInsertHandlerKernelBundlePtr(/*Insert=*/true);
     }
   }
 
@@ -1706,13 +1706,26 @@ class __SYCL_EXPORT handler {
   void setStateSpecConstSet();
   bool isStateExplicitKernelBundle() const;
 
+#ifndef __INTEL_PREVIEW_BREAKING_CHANGES
   std::shared_ptr<detail::kernel_bundle_impl>
   getOrInsertHandlerKernelBundle(bool Insert) const;
+#endif
+
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  // Rename to just getOrInsertHandlerKernelBundle
+#endif
+  detail::kernel_bundle_impl *
+  getOrInsertHandlerKernelBundlePtr(bool Insert) const;
 
   void setHandlerKernelBundle(kernel Kernel);
 
+#ifndef __INTEL_PREVIEW_BREAKING_CHANGES
   void setHandlerKernelBundle(
       const std::shared_ptr<detail::kernel_bundle_impl> &NewKernelBundleImpPtr);
+#endif
+
+  template <typename SharedPtrT>
+  void setHandlerKernelBundle(SharedPtrT &&NewKernelBundleImpPtr);
 
   void SetHostTask(std::function<void()> &&Func);
   void SetHostTask(std::function<void(interop_handle)> &&Func);
@@ -1760,6 +1773,8 @@ class __SYCL_EXPORT handler {
   /// called.
   void setUserFacingNodeType(ext::oneapi::experimental::node_type Type);
 
+  kernel_bundle<bundle_state::input> getKernelBundle() const;
+
 public:
   handler(const handler &) = delete;
   handler(handler &&) = delete;

sycl/include/sycl/kernel_bundle.hpp

@@ -1330,12 +1330,7 @@ void handler::set_specialization_constant(
 
   setStateSpecConstSet();
 
-  std::shared_ptr<detail::kernel_bundle_impl> KernelBundleImplPtr =
-      getOrInsertHandlerKernelBundle(/*Insert=*/true);
-
-  detail::createSyclObjFromImpl<kernel_bundle<bundle_state::input>>(
-      std::move(KernelBundleImplPtr))
-      .set_specialization_constant<SpecName>(Value);
+  getKernelBundle().set_specialization_constant<SpecName>(Value);
 }
 
 template <auto &SpecName>
@@ -1347,12 +1342,7 @@ handler::get_specialization_constant() const {
                           "Specialization constants cannot be read after "
                           "explicitly setting the used kernel bundle");
 
-  std::shared_ptr<detail::kernel_bundle_impl> KernelBundleImplPtr =
-      getOrInsertHandlerKernelBundle(/*Insert=*/true);
-
-  return detail::createSyclObjFromImpl<kernel_bundle<bundle_state::input>>(
-             std::move(KernelBundleImplPtr))
-      .get_specialization_constant<SpecName>();
+  return getKernelBundle().get_specialization_constant<SpecName>();
 }
 
 } // namespace _V1

sycl/source/CMakeLists.txt

@@ -269,6 +269,7 @@ set(SYCL_COMMON_SOURCES
     "detail/host_pipe_map.cpp"
     "detail/device_global_map.cpp"
     "detail/device_global_map_entry.cpp"
+    "detail/device_image_impl.cpp"
     "detail/device_impl.cpp"
     "detail/error_handling/error_handling.cpp"
     "detail/event_impl.cpp"

sycl/source/backend.cpp

@@ -306,7 +306,7 @@ make_kernel_bundle(ur_native_handle_t NativeHandle,
       ImageOriginInterop);
   device_image_plain DevImg{DevImgImpl};
 
-  return std::make_shared<kernel_bundle_impl>(TargetContext, Devices, DevImg);
+  return kernel_bundle_impl::create(TargetContext, Devices, DevImg);
 }
 
 // TODO: Unused. Remove when allowed.

sycl/source/detail/device_image_impl.cpp

@@ -0,0 +1,59 @@
+//==----------------- device_image_impl.cpp - SYCL device_image_impl -------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <detail/device_image_impl.hpp>
+#include <detail/kernel_bundle_impl.hpp>
+
+namespace sycl {
+inline namespace _V1 {
+namespace detail {
+
+std::shared_ptr<kernel_impl> device_image_impl::tryGetSourceBasedKernel(
+    std::string_view Name, const context &Context,
+    const kernel_bundle_impl &OwnerBundle,
+    const std::shared_ptr<device_image_impl> &Self) const {
+  if (!(getOriginMask() & ImageOriginKernelCompiler))
+    return nullptr;
+
+  assert(MRTCBinInfo);
+  std::string AdjustedName = adjustKernelName(Name);
+  if (MRTCBinInfo->MLanguage == syclex::source_language::sycl) {
+    auto &PM = ProgramManager::getInstance();
+    for (const std::string &Prefix : MRTCBinInfo->MPrefixes) {
+      auto KID = PM.tryGetSYCLKernelID(Prefix + AdjustedName);
+
+      if (!KID || !has_kernel(*KID))
+        continue;
+
+      auto UrProgram = get_ur_program_ref();
+      auto [UrKernel, CacheMutex, ArgMask] =
+          PM.getOrCreateKernel(Context, AdjustedName,
+                               /*PropList=*/{}, UrProgram);
+      return std::make_shared<kernel_impl>(UrKernel, *getSyclObjImpl(Context),
+                                           Self, OwnerBundle.shared_from_this(),
+                                           ArgMask, UrProgram, CacheMutex);
+    }
+    return nullptr;
+  }
+
+  ur_program_handle_t UrProgram = get_ur_program_ref();
+  const AdapterPtr &Adapter = getSyclObjImpl(Context)->getAdapter();
+  ur_kernel_handle_t UrKernel = nullptr;
+  Adapter->call<UrApiKind::urKernelCreate>(UrProgram, AdjustedName.c_str(),
+                                           &UrKernel);
+  // Kernel created by urKernelCreate is implicitly retained.
+
+  return std::make_shared<kernel_impl>(
+      UrKernel, *detail::getSyclObjImpl(Context), Self,
+      OwnerBundle.shared_from_this(), /*ArgMask=*/nullptr, UrProgram,
+      /*CacheMutex=*/nullptr);
+}
+
+} // namespace detail
+} // namespace _V1
+} // namespace sycl

sycl/source/detail/device_image_impl.hpp

@@ -617,45 +617,10 @@ class device_image_impl {
                MRTCBinInfo->MKernelNames.end();
   }
 
-  std::shared_ptr<kernel_impl> tryGetSourceBasedKernel(
-      std::string_view Name, const context &Context,
-      const std::shared_ptr<kernel_bundle_impl> &OwnerBundle,
-      const std::shared_ptr<device_image_impl> &Self) const {
-    if (!(getOriginMask() & ImageOriginKernelCompiler))
-      return nullptr;
-
-    assert(MRTCBinInfo);
-    std::string AdjustedName = adjustKernelName(Name);
-    if (MRTCBinInfo->MLanguage == syclex::source_language::sycl) {
-      auto &PM = ProgramManager::getInstance();
-      for (const std::string &Prefix : MRTCBinInfo->MPrefixes) {
-        auto KID = PM.tryGetSYCLKernelID(Prefix + AdjustedName);
-
-        if (!KID || !has_kernel(*KID))
-          continue;
-
-        auto UrProgram = get_ur_program_ref();
-        auto [UrKernel, CacheMutex, ArgMask] =
-            PM.getOrCreateKernel(Context, AdjustedName,
-                                 /*PropList=*/{}, UrProgram);
-        return std::make_shared<kernel_impl>(UrKernel, *getSyclObjImpl(Context),
-                                             Self, OwnerBundle, ArgMask,
-                                             UrProgram, CacheMutex);
-      }
-      return nullptr;
-    }
-
-    ur_program_handle_t UrProgram = get_ur_program_ref();
-    const AdapterPtr &Adapter = getSyclObjImpl(Context)->getAdapter();
-    ur_kernel_handle_t UrKernel = nullptr;
-    Adapter->call<UrApiKind::urKernelCreate>(UrProgram, AdjustedName.c_str(),
-                                             &UrKernel);
-    // Kernel created by urKernelCreate is implicitly retained.
-
-    return std::make_shared<kernel_impl>(
-        UrKernel, *detail::getSyclObjImpl(Context), Self, OwnerBundle,
-        /*ArgMask=*/nullptr, UrProgram, /*CacheMutex=*/nullptr);
-  }
+  std::shared_ptr<kernel_impl>
+  tryGetSourceBasedKernel(std::string_view Name, const context &Context,
+                          const kernel_bundle_impl &OwnerBundle,
+                          const std::shared_ptr<device_image_impl> &Self) const;
 
   bool hasDeviceGlobalName(const std::string &Name) const noexcept {
     if (!MRTCBinInfo.has_value())

sycl/source/detail/graph_impl.cpp

@@ -860,7 +860,7 @@ exec_graph_impl::enqueueNodeDirect(const sycl::context &Ctx,
     std::tie(CmdTraceEvent, InstanceID) = emitKernelInstrumentationData(
         StreamID, CGExec->MSyclKernel, CodeLoc, CGExec->MIsTopCodeLoc,
         CGExec->MKernelName.data(), CGExec->MKernelNameBasedCachePtr, nullptr,
-        CGExec->MNDRDesc, CGExec->MKernelBundle, CGExec->MArgs);
+        CGExec->MNDRDesc, CGExec->MKernelBundle.get(), CGExec->MArgs);
     if (CmdTraceEvent)
       sycl::detail::emitInstrumentationGeneral(
           StreamID, InstanceID, CmdTraceEvent, xpti::trace_task_begin, nullptr);
@@ -1536,8 +1536,7 @@ void exec_graph_impl::populateURKernelUpdateStructs(
     EliminatedArgMask = Kernel->getKernelArgMask();
   } else if (auto SyclKernelImpl =
                  KernelBundleImplPtr
-                     ? KernelBundleImplPtr->tryGetKernel(ExecCG.MKernelName,
-                                                         KernelBundleImplPtr)
+                     ? KernelBundleImplPtr->tryGetKernel(ExecCG.MKernelName)
                      : std::shared_ptr<kernel_impl>{nullptr}) {
     UrKernel = SyclKernelImpl->getHandleRef();
     EliminatedArgMask = SyclKernelImpl->getKernelArgMask();

sycl/source/detail/helpers.cpp

@@ -73,8 +73,7 @@ retrieveKernelBinary(queue_impl &Queue, KernelNameStrRefT KernelName,
     DeviceImage = KernelCG->MSyclKernel->getDeviceImage()->get_bin_image_ref();
     Program = KernelCG->MSyclKernel->getDeviceImage()->get_ur_program_ref();
   } else if (auto SyclKernelImpl =
-                 KernelBundleImpl ? KernelBundleImpl->tryGetKernel(
-                                        KernelName, KernelBundleImpl)
+                 KernelBundleImpl ? KernelBundleImpl->tryGetKernel(KernelName)
                                   : std::shared_ptr<kernel_impl>{nullptr}) {
     // Retrieve the device image from the kernel bundle.
     DeviceImage = SyclKernelImpl->getDeviceImage()->get_bin_image_ref();