[SYCL] Fix use-after-release kernel issue

After the changes in intel#17380, kernels looked up from a kernel bundle
could cause lifetime issues, as the lifetime of the returned objects had
to be handled differently based on the path taken. These changes move
back to only look up non-source-based kernels and use the stored kernel
objects otherwise.

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/source/detail/graph_impl.cpp

@@ -1474,11 +1474,10 @@ void exec_graph_impl::populateURKernelUpdateStructs(
   ur_kernel_handle_t UrKernel = nullptr;
   auto Kernel = ExecCG.MSyclKernel;
   auto KernelBundleImplPtr = ExecCG.MKernelBundle;
-  std::shared_ptr<sycl::detail::kernel_impl> SyclKernelImpl = nullptr;
   const sycl::detail::KernelArgMask *EliminatedArgMask = nullptr;
 
   if (auto SyclKernelImpl = KernelBundleImplPtr
-                                ? KernelBundleImplPtr->tryGetKernel(
+                                ? KernelBundleImplPtr->tryGetOfflineKernel(
                                       ExecCG.MKernelName, KernelBundleImplPtr)
                                 : std::shared_ptr<kernel_impl>{nullptr}) {
     UrKernel = SyclKernelImpl->getHandleRef();

sycl/source/detail/kernel_bundle_impl.hpp

@@ -537,7 +537,7 @@ class kernel_bundle_impl {
   get_kernel(const kernel_id &KernelID,
              const std::shared_ptr<detail::kernel_bundle_impl> &Self) const {
     if (std::shared_ptr<kernel_impl> KernelImpl =
-            tryGetOfflineKernel(KernelID, Self))
+            tryGetOfflineKernelImpl(KernelID, Self))
       return detail::createSyclObjFromImpl<kernel>(std::move(KernelImpl));
     throw sycl::exception(make_error_code(errc::invalid),
                           "The kernel bundle does not contain the kernel "
@@ -692,7 +692,7 @@ class kernel_bundle_impl {
     });
   }
 
-  std::shared_ptr<kernel_impl> tryGetOfflineKernel(
+  std::shared_ptr<kernel_impl> tryGetOfflineKernelImpl(
       const kernel_id &KernelID,
       const std::shared_ptr<detail::kernel_bundle_impl> &Self) const {
     using ImageImpl = std::shared_ptr<detail::device_image_impl>;
@@ -758,6 +758,17 @@ class kernel_bundle_impl {
         SelectedImage->get_ur_program_ref(), CacheMutex);
   }
 
+  std::shared_ptr<kernel_impl>
+  tryGetOfflineKernel(detail::KernelNameStrRefT Name,
+                      const std::shared_ptr<kernel_bundle_impl> &Self) const {
+    // Fall back to regular offline compiled kernel_bundle look-up.
+    if (std::optional<kernel_id> MaybeKernelID =
+            sycl::detail::ProgramManager::getInstance().tryGetSYCLKernelID(
+                Name))
+      return tryGetOfflineKernelImpl(*MaybeKernelID, Self);
+    return nullptr;
+  }
+
   std::shared_ptr<kernel_impl>
   tryGetKernel(detail::KernelNameStrRefT Name,
                const std::shared_ptr<kernel_bundle_impl> &Self) const {
@@ -774,11 +785,7 @@ class kernel_bundle_impl {
     }
 
     // Fall back to regular offline compiled kernel_bundle look-up.
-    if (std::optional<kernel_id> MaybeKernelID =
-            sycl::detail::ProgramManager::getInstance().tryGetSYCLKernelID(
-                Name))
-      return tryGetOfflineKernel(*MaybeKernelID, Self);
-    return nullptr;
+    return tryGetOfflineKernel(Name, Self);
   }
 
 private:

sycl/source/detail/scheduler/commands.cpp

@@ -2523,8 +2523,8 @@ getCGKernelInfo(const CGExecKernel &CommandGroup, ContextImplPtr ContextImpl,
 
   if (auto SyclKernelImpl =
           KernelBundleImplPtr
-              ? KernelBundleImplPtr->tryGetKernel(CommandGroup.MKernelName,
-                                                  KernelBundleImplPtr)
+              ? KernelBundleImplPtr->tryGetOfflineKernel(
+                    CommandGroup.MKernelName, KernelBundleImplPtr)
               : std::shared_ptr<kernel_impl>{nullptr}) {
     UrKernel = SyclKernelImpl->getHandleRef();
     DeviceImageImpl = SyclKernelImpl->getDeviceImage();

sycl/source/handler.cpp

@@ -445,10 +445,8 @@ event handler::finalize() {
       // Make sure implicit non-interop kernel bundles have the kernel
       if (!impl->isStateExplicitKernelBundle() &&
           !(MKernel && MKernel->isInterop()) &&
-          (KernelBundleImpPtr->empty() ||
-           KernelBundleImpPtr->hasSYCLOfflineImages()) &&
-          !KernelBundleImpPtr->tryGetKernel(MKernelName.data(),
-                                            KernelBundleImpPtr)) {
+          !KernelBundleImpPtr->tryGetOfflineKernel(MKernelName.data(),
+                                                   KernelBundleImpPtr)) {
         auto Dev =
             impl->MGraph ? impl->MGraph->getDevice() : MQueue->get_device();
         kernel_id KernelID =