[SYCL] Fix unreleased kernels obtained with program::get_kernel() (#1804)

Avoid unconditional retaining of internally created native kernels in
kernel_impl constructors, retain them in the caller when needed
instead.

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

Author: sergey-semenov

sycl/source/detail/kernel_impl.cpp

@@ -22,7 +22,12 @@ namespace detail {
 kernel_impl::kernel_impl(RT::PiKernel Kernel, ContextImplPtr Context)
     : kernel_impl(Kernel, Context,
                   std::make_shared<program_impl>(Context, Kernel),
-                  /*IsCreatedFromSource*/ true) {}
+                  /*IsCreatedFromSource*/ true) {
+  // This constructor is only called in the interoperability kernel constructor.
+  // Let the runtime caller handle native kernel retaining in other cases if
+  // it's needed.
+  getPlugin().call<PiApiKind::piKernelRetain>(MKernel);
+}
 
 kernel_impl::kernel_impl(RT::PiKernel Kernel, ContextImplPtr ContextImpl,
                          ProgramImplPtr ProgramImpl,
@@ -39,7 +44,6 @@ kernel_impl::kernel_impl(RT::PiKernel Kernel, ContextImplPtr ContextImpl,
     throw cl::sycl::invalid_parameter_error(
         "Input context must be the same as the context of cl_kernel",
         PI_INVALID_CONTEXT);
-  getPlugin().call<PiApiKind::piKernelRetain>(MKernel);
 }
 
 kernel_impl::kernel_impl(ContextImplPtr Context,

sycl/source/detail/program_impl.cpp

@@ -398,6 +398,7 @@ RT::PiKernel program_impl::get_pi_kernel(const string_class &KernelName) const {
   if (is_cacheable()) {
     Kernel = ProgramManager::getInstance().getOrCreateKernel(
         MProgramModuleHandle, get_context(), KernelName, this);
+    getPlugin().call<PiApiKind::piKernelRetain>(Kernel);
   } else {
     const detail::plugin &Plugin = getPlugin();
     RT::PiResult Err = Plugin.call_nocheck<PiApiKind::piKernelCreate>(

sycl/unittests/CMakeLists.txt

@@ -56,5 +56,6 @@ endfunction()
 
 add_subdirectory(misc)
 add_subdirectory(pi)
-add_subdirectory(thread_safety)
+add_subdirectory(program)
 add_subdirectory(scheduler)
+add_subdirectory(thread_safety)

sycl/unittests/program/CMakeLists.txt

@@ -0,0 +1,3 @@
+add_sycl_unittest(ProgramTests OBJECT
+  KernelRelease.cpp
+)

sycl/unittests/program/KernelRelease.cpp

@@ -0,0 +1,99 @@
+//==----------- KernelRelease.cpp --- kernel release unit test -------------==//
+//
+// 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 <CL/sycl.hpp>
+#include <detail/context_impl.hpp>
+#include <gtest/gtest.h>
+#include <helpers/PiMock.hpp>
+
+#include <iostream>
+#include <memory>
+
+using namespace cl::sycl;
+
+struct TestCtx {
+  TestCtx(context &Ctx) : Ctx{Ctx} {};
+
+  context &Ctx;
+  int KernelReferenceCount = 0;
+};
+
+std::unique_ptr<TestCtx> TestContext;
+
+pi_result redefinedProgramCreateWithSource(pi_context context, pi_uint32 count,
+                                           const char **strings,
+                                           const size_t *lengths,
+                                           pi_program *ret_program) {
+  return PI_SUCCESS;
+}
+
+pi_result
+redefinedProgramBuild(pi_program program, pi_uint32 num_devices,
+                      const pi_device *device_list, const char *options,
+                      void (*pfn_notify)(pi_program program, void *user_data),
+                      void *user_data) {
+  return PI_SUCCESS;
+}
+
+pi_result redefinedKernelCreate(pi_program program, const char *kernel_name,
+                                pi_kernel *ret_kernel) {
+  TestContext->KernelReferenceCount = 1;
+  return PI_SUCCESS;
+}
+
+pi_result redefinedKernelRetain(pi_kernel kernel) {
+  ++TestContext->KernelReferenceCount;
+  return PI_SUCCESS;
+}
+
+pi_result redefinedKernelRelease(pi_kernel kernel) {
+  --TestContext->KernelReferenceCount;
+  return PI_SUCCESS;
+}
+
+pi_result redefinedKernelGetInfo(pi_kernel kernel, pi_kernel_info param_name,
+                                 size_t param_value_size, void *param_value,
+                                 size_t *param_value_size_ret) {
+  EXPECT_EQ(param_name, PI_KERNEL_INFO_CONTEXT)
+      << "Unexpected kernel info requested";
+  auto *Result = reinterpret_cast<RT::PiContext *>(param_value);
+  RT::PiContext PiCtx =
+      detail::getSyclObjImpl(TestContext->Ctx)->getHandleRef();
+  *Result = PiCtx;
+  return PI_SUCCESS;
+}
+
+TEST(KernelReleaseTest, GetKernelRelease) {
+  unittest::PiMock Mock;
+  platform Plt = Mock.getPlatform();
+  if (Plt.is_host()) {
+    std::cerr << "The program/kernel methods are mostly no-op on the host "
+                 "device, the test is not run."
+              << std::endl;
+    return;
+  }
+
+  Mock.redefine<detail::PiApiKind::piclProgramCreateWithSource>(
+      redefinedProgramCreateWithSource);
+  Mock.redefine<detail::PiApiKind::piProgramBuild>(redefinedProgramBuild);
+  Mock.redefine<detail::PiApiKind::piKernelCreate>(redefinedKernelCreate);
+  Mock.redefine<detail::PiApiKind::piKernelRetain>(redefinedKernelRetain);
+  Mock.redefine<detail::PiApiKind::piKernelRelease>(redefinedKernelRelease);
+  Mock.redefine<detail::PiApiKind::piKernelGetInfo>(redefinedKernelGetInfo);
+
+  context Ctx{Plt};
+  TestContext.reset(new TestCtx(Ctx));
+
+  program Prg{Ctx};
+  Prg.build_with_source("");
+
+  { kernel Krnl = Prg.get_kernel(""); }
+
+  ASSERT_EQ(TestContext->KernelReferenceCount, 0)
+      << "Reference count not equal to 0 after kernel destruction";
+}