From 1778ad5557ec5bdd1a0e3cce44cc33b337bf3104 Mon Sep 17 00:00:00 2001
From: "Larsen, Steffen" <steffen.larsen@intel.com>
Date: Tue, 22 Apr 2025 21:54:30 -0700
Subject: [PATCH] [SYCL] Fix use-after-release kernel issue

After the changes in #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 <steffen.larsen@intel.com>
---
 sycl/source/detail/graph_impl.cpp         |  3 +--
 sycl/source/detail/kernel_bundle_impl.hpp | 21 ++++++++++++++-------
 sycl/source/detail/scheduler/commands.cpp |  4 ++--
 sycl/source/handler.cpp                   |  6 ++----
 4 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/sycl/source/detail/graph_impl.cpp b/sycl/source/detail/graph_impl.cpp
index b5f1f53e32d7f..42f26547a58df 100644
--- a/sycl/source/detail/graph_impl.cpp
+++ b/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();
diff --git a/sycl/source/detail/kernel_bundle_impl.hpp b/sycl/source/detail/kernel_bundle_impl.hpp
index 12ac9a47b7c07..183215c5aac0a 100644
--- a/sycl/source/detail/kernel_bundle_impl.hpp
+++ b/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:
diff --git a/sycl/source/detail/scheduler/commands.cpp b/sycl/source/detail/scheduler/commands.cpp
index b8b07e0053881..9a447c32a98b3 100644
--- a/sycl/source/detail/scheduler/commands.cpp
+++ b/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();
diff --git a/sycl/source/handler.cpp b/sycl/source/handler.cpp
index 97e6e316fea72..60ee539e8fb70 100644
--- a/sycl/source/handler.cpp
+++ b/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 =