[SYCL] Fix kernel program cache for multiple devices and refactor some unit tests (#5017)

When multiple devices are used in the same context, the problem of piProgram double free is arised.
In program_manager::build() created program is shared with multiple devices and is cached in KernelProgramCache as many times as there are devices. When cache is destructed, the same program will be released for every device, but piProgramRetain was called only for one device, which leads to double free. Adding piProgramRetain for each device solves this problem.
Deleting kernel from cache after its release prevents double free in the case a kernel was created for a program that shared with multiple devices.
Unittest with emulating multiple gpu is added.
Signed-off-by: mdimakov <[email protected]>

Author: maximdimakov

sycl/source/detail/kernel_program_cache.cpp

@@ -33,6 +33,7 @@ KernelProgramCache::~KernelProgramCache() {
           Plugin.call<PiApiKind::piKernelRelease>(Kern);
         }
       }
+      MKernelsPerProgramCache.erase(KernIt);
     }
 
     const detail::plugin &Plugin = MParentContext->getPlugin();

sycl/source/detail/program_manager/program_manager.cpp

@@ -1813,7 +1813,11 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
   // Cache supports key with once device only, but here we have multiple
   // devices a program is built for, so add the program to the cache for all
   // other devices.
-  auto CacheOtherDevices = [ResProgram]() { return ResProgram; };
+  const detail::plugin &Plugin = ContextImpl->getPlugin();
+  auto CacheOtherDevices = [ResProgram, &Plugin]() {
+    Plugin.call<PiApiKind::piProgramRetain>(ResProgram);
+    return ResProgram;
+  };
 
   // The program for device "0" is already added to the cache during the first
   // call to getOrBuild, so starting with "1"
@@ -1833,7 +1837,6 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
   // devive_image_impl shares ownership of PIProgram with, at least, program
   // cache. The ref counter will be descremented in the destructor of
   // device_image_impl
-  const detail::plugin &Plugin = ContextImpl->getPlugin();
   Plugin.call<PiApiKind::piProgramRetain>(ResProgram);
 
   DeviceImageImplPtr ExecImpl = std::make_shared<detail::device_image_impl>(

sycl/unittests/helpers/CommonRedefinitions.hpp

@@ -6,7 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include <CL/sycl.hpp>
 #include <helpers/PiImage.hpp>
 #include <helpers/PiMock.hpp>
 

sycl/unittests/helpers/PiImage.hpp

@@ -8,7 +8,6 @@
 
 #pragma once
 
-#include <CL/sycl.hpp>
 #include <CL/sycl/detail/common.hpp>
 #include <CL/sycl/detail/pi.hpp>
 #include <detail/platform_impl.hpp>

sycl/unittests/helpers/PiMock.hpp

@@ -27,9 +27,11 @@
 
 #pragma once
 
-#include <CL/sycl.hpp>
 #include <CL/sycl/detail/common.hpp>
 #include <CL/sycl/detail/pi.hpp>
+#include <CL/sycl/device_selector.hpp>
+#include <CL/sycl/platform.hpp>
+#include <CL/sycl/queue.hpp>
 #include <detail/platform_impl.hpp>
 
 #include <functional>

sycl/unittests/helpers/TestKernel.hpp

@@ -0,0 +1,55 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "PiImage.hpp"
+
+class TestKernel;
+
+__SYCL_INLINE_NAMESPACE(cl) {
+namespace sycl {
+namespace detail {
+template <> struct KernelInfo<TestKernel> {
+  static constexpr unsigned getNumParams() { return 0; }
+  static const kernel_param_desc_t &getParamDesc(int) {
+    static kernel_param_desc_t Dummy;
+    return Dummy;
+  }
+  static constexpr const char *getName() { return "TestKernel"; }
+  static constexpr bool isESIMD() { return false; }
+  static constexpr bool callsThisItem() { return false; }
+  static constexpr bool callsAnyThisFreeFunction() { return false; }
+};
+
+} // namespace detail
+} // namespace sycl
+} // __SYCL_INLINE_NAMESPACE(cl)
+
+static sycl::unittest::PiImage generateDefaultImage() {
+  using namespace sycl::unittest;
+
+  PiPropertySet PropSet;
+
+  std::vector<unsigned char> Bin{0, 1, 2, 3, 4, 5}; // Random data
+
+  PiArray<PiOffloadEntry> Entries = makeEmptyKernels({"TestKernel"});
+
+  PiImage Img{PI_DEVICE_BINARY_TYPE_SPIRV,            // Format
+              __SYCL_PI_DEVICE_BINARY_TARGET_SPIRV64, // DeviceTargetSpec
+              "",                                     // Compile options
+              "",                                     // Link options
+              std::move(Bin),
+              std::move(Entries),
+              std::move(PropSet)};
+
+  return Img;
+}
+
+static sycl::unittest::PiImage Img = generateDefaultImage();
+static sycl::unittest::PiImageArray<1> ImgArray{&Img};

sycl/unittests/kernel-and-program/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_sycl_unittest(KernelAndProgramTests OBJECT
   Cache.cpp
+  MultipleDevsCache.cpp
   KernelRelease.cpp
   KernelInfo.cpp
   DeviceInfo.cpp

sycl/unittests/kernel-and-program/Cache.cpp

@@ -51,11 +51,9 @@ struct MockKernelInfo {
 template <> struct KernelInfo<TestKernel> : public MockKernelInfo {
   static constexpr const char *getName() { return "TestKernel"; }
 };
-
 template <> struct KernelInfo<TestKernel2> : public MockKernelInfo {
   static constexpr const char *getName() { return "TestKernel2"; }
 };
-
 } // namespace detail
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)
@@ -117,6 +115,7 @@ static pi_result redefinedKernelCreate(pi_program program,
                                        pi_kernel *ret_kernel) {
   return PI_SUCCESS;
 }
+
 static pi_result redefinedKernelRelease(pi_kernel kernel) { return PI_SUCCESS; }
 
 class KernelAndProgramCacheTest : public ::testing::Test {

sycl/unittests/kernel-and-program/MultipleDevsCache.cpp

@@ -0,0 +1,184 @@
+//==--- KPCache.cpp --- KernelProgramCache for multiple devices 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
+//
+//===----------------------------------------------------------------------===//
+
+#define SYCL2020_DISABLE_DEPRECATION_WARNINGS
+
+#include "detail/context_impl.hpp"
+#include "detail/kernel_bundle_impl.hpp"
+#include "detail/kernel_program_cache.hpp"
+#include <helpers/CommonRedefinitions.hpp>
+#include <helpers/PiImage.hpp>
+#include <helpers/PiMock.hpp>
+#include <helpers/TestKernel.hpp>
+
+#include <gtest/gtest.h>
+
+#include <iostream>
+
+using namespace sycl;
+
+static pi_result redefinedContextCreate(
+    const pi_context_properties *properties, pi_uint32 num_devices,
+    const pi_device *devices,
+    void (*pfn_notify)(const char *errinfo, const void *private_info, size_t cb,
+                       void *user_data),
+    void *user_data, pi_context *ret_context) {
+  *ret_context = reinterpret_cast<pi_context>(123);
+  return PI_SUCCESS;
+}
+
+static pi_result redefinedContextRelease(pi_context context) {
+  return PI_SUCCESS;
+}
+
+static pi_result redefinedDevicesGet(pi_platform platform,
+                                     pi_device_type device_type,
+                                     pi_uint32 num_entries, pi_device *devices,
+                                     pi_uint32 *num_devices) {
+  if (num_devices) {
+    *num_devices = static_cast<pi_uint32>(2);
+    return PI_SUCCESS;
+  }
+
+  if (num_entries == 2 && devices) {
+    devices[0] = reinterpret_cast<pi_device>(1111);
+    devices[1] = reinterpret_cast<pi_device>(2222);
+  }
+  return PI_SUCCESS;
+}
+
+static pi_result redefinedDeviceGetInfo(pi_device device,
+                                        pi_device_info param_name,
+                                        size_t param_value_size,
+                                        void *param_value,
+                                        size_t *param_value_size_ret) {
+  if (param_name == PI_DEVICE_INFO_TYPE) {
+    auto *Result = reinterpret_cast<_pi_device_type *>(param_value);
+    *Result = PI_DEVICE_TYPE_GPU;
+  }
+  if (param_name == PI_DEVICE_INFO_COMPILER_AVAILABLE) {
+    auto *Result = reinterpret_cast<pi_bool *>(param_value);
+    *Result = true;
+  }
+  return PI_SUCCESS;
+}
+
+static pi_result redefinedDeviceRetain(pi_device device) { return PI_SUCCESS; }
+
+static pi_result redefinedDeviceRelease(pi_device device) { return PI_SUCCESS; }
+
+static pi_result redefinedQueueCreate(pi_context context, pi_device device,
+                                      pi_queue_properties properties,
+                                      pi_queue *queue) {
+  *queue = reinterpret_cast<pi_queue>(1234);
+  return PI_SUCCESS;
+}
+
+static pi_result redefinedQueueRelease(pi_queue command_queue) {
+  return PI_SUCCESS;
+}
+
+static size_t ProgramNum = 12345;
+static pi_result redefinedProgramCreate(pi_context context, const void *il,
+                                        size_t length,
+                                        pi_program *res_program) {
+  size_t CurrentProgram = ProgramNum;
+  *res_program = reinterpret_cast<pi_program>(CurrentProgram);
+  ++ProgramNum;
+  return PI_SUCCESS;
+}
+
+static int RetainCounter = 0;
+static pi_result redefinedProgramRetain(pi_program program) {
+  ++RetainCounter;
+  return PI_SUCCESS;
+}
+
+static int KernelReleaseCounter = 0;
+static pi_result redefinedKernelRelease(pi_kernel kernel) {
+  ++KernelReleaseCounter;
+  return PI_SUCCESS;
+}
+
+class MultipleDeviceCacheTest : public ::testing::Test {
+public:
+  MultipleDeviceCacheTest() : Plt{default_selector()} {}
+
+protected:
+  void SetUp() override {
+    if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+      return;
+    }
+
+    Mock = std::make_unique<unittest::PiMock>(Plt);
+
+    setupDefaultMockAPIs(*Mock);
+    Mock->redefine<detail::PiApiKind::piDevicesGet>(redefinedDevicesGet);
+    Mock->redefine<detail::PiApiKind::piDeviceGetInfo>(redefinedDeviceGetInfo);
+    Mock->redefine<detail::PiApiKind::piDeviceRetain>(redefinedDeviceRetain);
+    Mock->redefine<detail::PiApiKind::piDeviceRelease>(redefinedDeviceRelease);
+    Mock->redefine<detail::PiApiKind::piContextCreate>(redefinedContextCreate);
+    Mock->redefine<detail::PiApiKind::piContextRelease>(
+        redefinedContextRelease);
+    Mock->redefine<detail::PiApiKind::piQueueCreate>(redefinedQueueCreate);
+    Mock->redefine<detail::PiApiKind::piQueueRelease>(redefinedQueueRelease);
+    Mock->redefine<detail::PiApiKind::piProgramRetain>(redefinedProgramRetain);
+    Mock->redefine<detail::PiApiKind::piProgramCreate>(redefinedProgramCreate);
+    Mock->redefine<detail::PiApiKind::piKernelRelease>(redefinedKernelRelease);
+  }
+
+protected:
+  std::unique_ptr<unittest::PiMock> Mock;
+  platform Plt;
+};
+
+// Test that program is retained for each device and each kernel is released
+// once
+TEST_F(MultipleDeviceCacheTest, ProgramRetain) {
+  if (Plt.is_host() || Plt.get_backend() != backend::opencl) {
+    return;
+  }
+  {
+    std::vector<sycl::device> Devices = Plt.get_devices(info::device_type::gpu);
+    sycl::context Context(Devices);
+    sycl::queue Queue(Context, Devices[0]);
+    assert(Devices.size() == 2);
+
+    auto Bundle = cl::sycl::get_kernel_bundle<sycl::bundle_state::input>(
+        Queue.get_context());
+    Queue.submit(
+        [&](cl::sycl::handler &cgh) { cgh.single_task<TestKernel>([]() {}); });
+
+    auto BundleObject = cl::sycl::build(Bundle, Bundle.get_devices());
+    auto KernelID = cl::sycl::get_kernel_id<TestKernel>();
+    auto Kernel = BundleObject.get_kernel(KernelID);
+
+    // Because of emulating 2 devices program is retained for each one in
+    // build(). It is also depends on number of device images. This test has one
+    // image, but other tests can create other images. Additional variable is
+    // added to control count of piProgramRetain calls
+    auto BundleImpl = getSyclObjImpl(Bundle);
+    int NumRetains = BundleImpl->size() * 2;
+
+    EXPECT_EQ(RetainCounter, NumRetains)
+        << "Expect " << NumRetains << " piProgramRetain calls";
+
+    auto CtxImpl = detail::getSyclObjImpl(Context);
+    detail::KernelProgramCache::KernelCacheT &KernelCache =
+        CtxImpl->getKernelProgramCache().acquireKernelsPerProgramCache().get();
+
+    EXPECT_EQ(KernelCache.size(), (size_t)2) << "Expect 2 kernels in cache";
+  }
+  // First kernel creating is called in handler::single_task().
+  // kernel_bundle::get_kernel() creates a kernel and shares it with created
+  // programs. Also the kernel is retained in kernel_bundle::get_kernel(). A
+  // kernel is removed from cache if piKernelRelease was called for it, so it
+  // will not be removed twice for the other programs. As a result we must
+  // expect 3 piKernelRelease calls.
+  EXPECT_EQ(KernelReleaseCounter, 3) << "Expect 3 piKernelRelease calls";
+}

sycl/unittests/pi/PiMock.cpp

@@ -6,10 +6,12 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include <gtest/gtest.h>
 #include <helpers/PiMock.hpp>
+
 #include <detail/queue_impl.hpp>
 
+#include <gtest/gtest.h>
+
 using namespace cl::sycl;
 
 pi_result piProgramBuildRedefine(pi_program, pi_uint32, const pi_device *,