[SYCL] Fix issue of acquring kernel twice (#11953)

In #11751, ref counting of kernels objects was changed to be more
accurate in order to allow for in-memory caching to be disabled. When
getting a kernel form the cache, the ref count the kernel handle is now
incremented (when caching is enabled). Thus, a method like
`ProgramManager::getOrCreateKernel` will increment the ref count of the
kernel it gets. However, in `enqueueImpKernel`, when enqueuing a kernel
with a kernel bundle,`ProgramManager::getOrCreateKernel` is called
twice, first indirectly by:


https://github.com/intel/llvm/blob/c43a90f28eebfcdf1bc1d55430485e2834790a60/sycl/source/detail/scheduler/commands.cpp#L2527-L2528

and second directly by:


https://github.com/intel/llvm/blob/c43a90f28eebfcdf1bc1d55430485e2834790a60/sycl/source/detail/scheduler/commands.cpp#L2538-L2548

This means that the ref count of the acquired kernel is incremented
twice, yet the rest of the function will only free once, which leads to
a leak of the kernel. As the second comment and asserts say, the only
need for the second call to `getOrCreateKernel` is to fetch the mutex
associated to the cached kernel retrieved from the first call, so this
PR adjusts `get_kernel` to save this mutex and forgo this extra
`getOrCreateKernel` call and unintentional additional ref count.

Author: jzc

sycl/source/detail/kernel_bundle_impl.hpp

@@ -531,15 +531,14 @@ class kernel_bundle_impl {
                             "The kernel bundle does not contain the kernel "
                             "identified by kernelId.");
 
-    sycl::detail::pi::PiKernel Kernel = nullptr;
-    const KernelArgMask *ArgMask = nullptr;
-    std::tie(Kernel, std::ignore, ArgMask) =
+    auto [Kernel, CacheMutex, ArgMask] =
         detail::ProgramManager::getInstance().getOrCreateKernel(
             MContext, KernelID.get_name(), /*PropList=*/{},
             SelectedImage->get_program_ref());
 
-    std::shared_ptr<kernel_impl> KernelImpl = std::make_shared<kernel_impl>(
-        Kernel, detail::getSyclObjImpl(MContext), SelectedImage, Self, ArgMask);
+    std::shared_ptr<kernel_impl> KernelImpl =
+        std::make_shared<kernel_impl>(Kernel, detail::getSyclObjImpl(MContext),
+                                      SelectedImage, Self, ArgMask, CacheMutex);
 
     return detail::createSyclObjFromImpl<kernel>(KernelImpl);
   }

sycl/source/detail/kernel_impl.cpp

@@ -61,11 +61,11 @@ kernel_impl::kernel_impl(sycl::detail::pi::PiKernel Kernel,
                          ContextImplPtr ContextImpl,
                          DeviceImageImplPtr DeviceImageImpl,
                          KernelBundleImplPtr KernelBundleImpl,
-                         const KernelArgMask *ArgMask)
+                         const KernelArgMask *ArgMask, std::mutex *CacheMutex)
     : MKernel(Kernel), MContext(std::move(ContextImpl)), MProgramImpl(nullptr),
       MCreatedFromSource(false), MDeviceImageImpl(std::move(DeviceImageImpl)),
       MKernelBundleImpl(std::move(KernelBundleImpl)),
-      MKernelArgMaskPtr{ArgMask} {
+      MKernelArgMaskPtr{ArgMask}, MCacheMutex{CacheMutex} {
   MIsInterop = MKernelBundleImpl->isInterop();
 }
 

sycl/source/detail/kernel_impl.hpp

@@ -73,7 +73,7 @@ class kernel_impl {
   kernel_impl(sycl::detail::pi::PiKernel Kernel, ContextImplPtr ContextImpl,
               DeviceImageImplPtr DeviceImageImpl,
               KernelBundleImplPtr KernelBundleImpl,
-              const KernelArgMask *ArgMask);
+              const KernelArgMask *ArgMask, std::mutex *CacheMutex);
 
   /// Constructs a SYCL kernel for host device
   ///
@@ -183,6 +183,7 @@ class kernel_impl {
   }
 
   const KernelArgMask *getKernelArgMask() const { return MKernelArgMaskPtr; }
+  std::mutex *getCacheMutex() const { return MCacheMutex; }
 
 private:
   sycl::detail::pi::PiKernel MKernel;
@@ -194,6 +195,7 @@ class kernel_impl {
   bool MIsInterop = false;
   std::mutex MNoncacheableEnqueueMutex;
   const KernelArgMask *MKernelArgMaskPtr;
+  std::mutex *MCacheMutex;
 
   bool isBuiltInKernel(const device &Device) const;
   void checkIfValidForNumArgsInfoQuery() const;

sycl/source/detail/program_manager/program_manager.cpp

@@ -622,6 +622,8 @@ static void emitBuiltProgramInfo(const pi_program &Prog,
   }
 }
 
+// When caching is enabled, the returned PiProgram will already have
+// its ref count incremented.
 sycl::detail::pi::PiProgram ProgramManager::getBuiltPIProgram(
     const ContextImplPtr &ContextImpl, const DeviceImplPtr &DeviceImpl,
     const std::string &KernelName, bool JITCompilationIsRequired) {
@@ -739,6 +741,8 @@ sycl::detail::pi::PiProgram ProgramManager::getBuiltPIProgram(
   return *BuildResult->Ptr.load();
 }
 
+// When caching is enabled, the returned PiProgram and PiKernel will
+// already have their ref count incremented.
 std::tuple<sycl::detail::pi::PiKernel, std::mutex *, const KernelArgMask *,
            sycl::detail::pi::PiProgram>
 ProgramManager::getOrCreateKernel(const ContextImplPtr &ContextImpl,
@@ -2432,6 +2436,8 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
   return createSyclObjFromImpl<device_image_plain>(ExecImpl);
 }
 
+// When caching is enabled, the returned PiKernel will already have
+// its ref count incremented.
 std::tuple<sycl::detail::pi::PiKernel, std::mutex *, const KernelArgMask *>
 ProgramManager::getOrCreateKernel(const context &Context,
                                   const std::string &KernelName,

sycl/source/detail/scheduler/commands.cpp

@@ -2535,18 +2535,7 @@ pi_int32 enqueueImpKernel(
     Program = DeviceImageImpl->get_program_ref();
 
     EliminatedArgMask = SyclKernelImpl->getKernelArgMask();
-    // When caching is enabled, kernel objects can be shared,
-    // so we need to retrieve the mutex associated to it via
-    // getOrCreateKernel
-    if (SYCLConfig<SYCL_CACHE_IN_MEM>::get()) {
-      auto [CachedKernel, CachedKernelMutex, CachedEliminatedArgMask] =
-          detail::ProgramManager::getInstance().getOrCreateKernel(
-              KernelBundleImplPtr->get_context(), KernelName,
-              /*PropList=*/{}, Program);
-      assert(CachedKernel == Kernel);
-      assert(CachedEliminatedArgMask == EliminatedArgMask);
-      KernelMutex = CachedKernelMutex;
-    }
+    KernelMutex = SyclKernelImpl->getCacheMutex();
   } else if (nullptr != MSyclKernel) {
     assert(MSyclKernel->get_info<info::kernel::context>() ==
            Queue->get_context());

sycl/test-e2e/KernelAndProgram/disable-caching.cpp

@@ -2,9 +2,9 @@
 // if and only if caching is disabled.
 
 // RUN: %{build} -o %t.out
-// RUN: env SYCL_PI_TRACE=-1 SYCL_CACHE_IN_MEM=0 %{run} %t.out \
+// RUN: env ZE_DEBUG=-6 SYCL_PI_TRACE=-1 SYCL_CACHE_IN_MEM=0 %{run} %t.out \
 // RUN: | FileCheck %s
-// RUN: env SYCL_PI_TRACE=-1 %{run} %t.out \
+// RUN: env ZE_DEBUG=-6 SYCL_PI_TRACE=-1 %{run} %t.out \
 // RUN: | FileCheck %s --check-prefixes=CHECK-CACHE
 #include <sycl/sycl.hpp>
 
@@ -25,8 +25,10 @@ int main() {
 
   // CHECK-CACHE: piProgramCreate
   // CHECK-CACHE: piProgramRetain
+  // CHECK-CACHE-NOT: piProgramRetain
   // CHECK-CACHE: piKernelCreate
   // CHECK-CACHE: piKernelRetain
+  // CHECK-CACHE-NOT: piKernelCreate
   // CHECK-CACHE: piEnqueueKernelLaunch
   // CHECK-CACHE: piKernelRelease
   // CHECK-CACHE: piProgramRelease
@@ -44,8 +46,10 @@ int main() {
 
   // CHECK-CACHE: piProgramCreate
   // CHECK-CACHE: piProgramRetain
+  // CHECK-CACHE-NOT: piProgramRetain
   // CHECK-CACHE: piKernelCreate
   // CHECK-CACHE: piKernelRetain
+  // CHECK-CACHE-NOT: piKernelCreate
   // CHECK-CACHE: piEnqueueKernelLaunch
   // CHECK-CACHE: piKernelRelease
   // CHECK-CACHE: piProgramRelease
@@ -62,8 +66,10 @@ int main() {
 
   // CHECK-CACHE: piProgramCreate
   // CHECK-CACHE: piProgramRetain
+  // CHECK-CACHE-NOT: piProgramRetain
   // CHECK-CACHE: piKernelCreate
   // CHECK-CACHE: piKernelRetain
+  // CHECK-CACHE-NOT: piKernelCreate
   // CHECK-CACHE: piEnqueueKernelLaunch
   // CHECK-CACHE: piKernelRelease
   // CHECK-CACHE: piProgramRelease