[SYCL] ProgramManager: Remove invalidated entries in NativePrograms before insert (#15973)

Fixes #14972

This commit erases existing entries with ur handle that is just
created/returned by backend. All existing entries in NativePrograms are
known to be invalid in this case.
Could not erase them on UrProgramRelease call since we have no tracking
of program handle references on SYCL RT level and it is not feasible to
add it. Obtaining ref count from ur is not thread safe and not a feature
to base product on.

---------

Signed-off-by: Tikhomirova, Kseniya <[email protected]>
Co-authored-by: Sergey Semenov <[email protected]>

Author: KseniyaTikhomirova

sycl/source/detail/program_manager/program_manager.cpp

@@ -509,7 +509,7 @@ std::pair<ur_program_handle_t, bool> ProgramManager::getOrCreateURProgram(
     const std::vector<const RTDeviceBinaryImage *> &AllImages,
     const context &Context, const std::vector<device> &Devices,
     const std::string &CompileAndLinkOptions, SerializedObj SpecConsts) {
-  ur_program_handle_t NativePrg; // TODO: Or native?
+  ur_program_handle_t NativePrg;
 
   // Get binaries for each device (1:1 correpsondence with input Devices).
   auto Binaries = PersistentDeviceCodeCache::getItemFromDisc(
@@ -768,7 +768,8 @@ setSpecializationConstants(const std::shared_ptr<device_image_impl> &InputImpl,
   }
 }
 
-static inline void CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) {
+static inline void
+CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) {
 #ifndef SYCL_RT_ZSTD_NOT_AVAIABLE
   if (auto CompImg = dynamic_cast<CompressedRTDeviceBinaryImage *>(Img))
     if (CompImg->IsCompressed())
@@ -913,6 +914,11 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
 
     {
       std::lock_guard<std::mutex> Lock(MNativeProgramsMutex);
+      // NativePrograms map does not intend to keep reference to program handle,
+      // so keys in the map can be invalid (reference count went to zero and the
+      // underlying program disposed of). Protecting from incorrect values by
+      // removal of map entries with same handle (obviously invalid entries).
+      std::ignore = NativePrograms.erase(BuiltProgram.get());
       for (const RTDeviceBinaryImage *Img : ImgWithDeps) {
         NativePrograms.insert({BuiltProgram.get(), Img});
       }
@@ -2747,6 +2753,11 @@ ProgramManager::link(const DevImgPlainWithDeps &ImgWithDeps,
 
   {
     std::lock_guard<std::mutex> Lock(MNativeProgramsMutex);
+    // NativePrograms map does not intend to keep reference to program handle,
+    // so keys in the map can be invalid (reference count went to zero and the
+    // underlying program disposed of). Protecting from incorrect values by
+    // removal of map entries with same handle (obviously invalid entries).
+    std::ignore = NativePrograms.erase(LinkedProg);
     for (const device_image_plain &Img : ImgWithDeps) {
       NativePrograms.insert(
           {LinkedProg, getSyclObjImpl(Img)->get_bin_image_ref()});

sycl/source/detail/program_manager/program_manager.hpp

@@ -46,6 +46,9 @@ extern "C" __SYCL_EXPORT void __sycl_unregister_lib(sycl_device_binaries desc);
 
 // +++ }
 
+// For testing purposes
+class ProgramManagerTest;
+
 namespace sycl {
 inline namespace _V1 {
 class context;
@@ -494,6 +497,8 @@ class ProgramManager {
   using MaterializedEntries =
       std::map<std::vector<unsigned char>, ur_kernel_handle_t>;
   std::unordered_map<std::string, MaterializedEntries> m_MaterializedKernels;
+
+  friend class ::ProgramManagerTest;
 };
 } // namespace detail
 } // namespace _V1

sycl/unittests/program_manager/arg_mask/EliminatedArgMask.cpp

@@ -6,14 +6,17 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <detail/config.hpp>
 #include <detail/handler_impl.hpp>
 #include <detail/kernel_bundle_impl.hpp>
+#include <detail/program_manager/program_manager.hpp>
 #include <detail/queue_impl.hpp>
 #include <detail/scheduler/commands.hpp>
 #include <sycl/sycl.hpp>
 
 #include <helpers/MockDeviceImage.hpp>
 #include <helpers/MockKernelInfo.hpp>
+#include <helpers/ScopedEnvVar.hpp>
 #include <helpers/UrMock.hpp>
 
 #include <gtest/gtest.h>
@@ -98,28 +101,17 @@ static sycl::unittest::MockDeviceImageArray<1> EAMImgArray{&EAMImg};
 static sycl::unittest::MockDeviceImageArray<1> EAM2ImgArray{&EAM2Img};
 static sycl::unittest::MockDeviceImageArray<1> EAM3ImgArray{&EAM3Img};
 
-// ur_program_handle_t address is used as a key for ProgramManager::NativePrograms
-// storage. redefinedProgramLinkCommon makes ur_program_handle_t address equal to 0x1.
-// Make sure that size of Bin is different for device images used in these tests
-// and greater than 1.
+// ur_program_handle_t address is used as a key for
+// ProgramManager::NativePrograms storage. redefinedProgramLinkCommon makes
+// ur_program_handle_t address equal to 0x1. Make sure that size of Bin is
+// different for device images used in these tests and greater than 1.
 inline ur_result_t redefinedProgramCreateEAM(void *pParams) {
   auto params = *static_cast<ur_program_create_with_il_params_t *>(pParams);
   static size_t UrProgramAddr = 2;
   **params.pphProgram = reinterpret_cast<ur_program_handle_t>(UrProgramAddr++);
   return UR_RESULT_SUCCESS;
 }
 
-mock::dummy_handle_t_ FixedHandle;
-inline ur_result_t setFixedProgramPtr(void *pParams) {
-  auto params = *static_cast<ur_program_create_with_il_params_t *>(pParams);
-  **params.pphProgram = reinterpret_cast<ur_program_handle_t>(&FixedHandle);
-  return UR_RESULT_SUCCESS;
-}
-inline ur_result_t releaseFixedProgramPtr(void *pParams) {
-  // Do nothing
-  return UR_RESULT_SUCCESS;
-}
-
 class MockHandler : public sycl::handler {
 
 public:
@@ -218,6 +210,53 @@ TEST(EliminatedArgMask, KernelBundleWith2Kernels) {
   EXPECT_EQ(*EliminatedArgMask, ExpElimArgMask);
 }
 
+std::vector<std::unique_ptr<mock::dummy_handle_t_>> UsedProgramHandles;
+std::vector<std::unique_ptr<mock::dummy_handle_t_>> ProgramHandlesToReuse;
+inline ur_result_t setFixedProgramPtr(void *pParams) {
+  auto params = *static_cast<ur_program_create_with_il_params_t *>(pParams);
+  if (!ProgramHandlesToReuse.empty()) {
+    auto it = ProgramHandlesToReuse.begin() + 1;
+    std::move(ProgramHandlesToReuse.begin(), it,
+              std::back_inserter(UsedProgramHandles));
+    ProgramHandlesToReuse.erase(ProgramHandlesToReuse.begin(), it);
+  } else
+    UsedProgramHandles.push_back(
+        std::make_unique<mock::dummy_handle_t_>(sizeof(unsigned)));
+  **params.pphProgram =
+      reinterpret_cast<ur_program_handle_t>(UsedProgramHandles.back().get());
+  return UR_RESULT_SUCCESS;
+}
+inline ur_result_t releaseFixedProgramPtr(void *pParams) {
+  auto params = *static_cast<ur_program_release_params_t *>(pParams);
+  {
+    auto it = std::find_if(
+        UsedProgramHandles.begin(), UsedProgramHandles.end(),
+        [&params](const std::unique_ptr<mock::dummy_handle_t_> &item) {
+          return reinterpret_cast<ur_program_handle_t>(item.get()) ==
+                 *params.phProgram;
+        });
+    if (it == UsedProgramHandles.end())
+      return UR_RESULT_SUCCESS;
+    std::move(it, it + 1, std::back_inserter(ProgramHandlesToReuse));
+    UsedProgramHandles.erase(it, it + 1);
+  }
+  return UR_RESULT_SUCCESS;
+}
+
+inline ur_result_t customProgramRetain(void *pParams) {
+  // do nothing
+  return UR_RESULT_SUCCESS;
+}
+
+class ProgramManagerTest {
+public:
+  static std::unordered_multimap<ur_program_handle_t,
+                                 const sycl::detail::RTDeviceBinaryImage *> &
+  getNativePrograms() {
+    return sycl::detail::ProgramManager::getInstance().NativePrograms;
+  }
+};
+
 // It's possible for the same handle to be reused for multiple distinct programs
 // This can happen if a program is released (freeing underlying memory) and then
 // a new program happens to get given that same memory for its handle.
@@ -227,6 +266,7 @@ TEST(EliminatedArgMask, KernelBundleWith2Kernels) {
 TEST(EliminatedArgMask, ReuseOfHandleValues) {
   sycl::detail::ProgramManager &PM =
       sycl::detail::ProgramManager::getInstance();
+  auto &NativePrograms = ProgramManagerTest::getNativePrograms();
 
   ur_program_handle_t ProgBefore = nullptr;
   ur_program_handle_t ProgAfter = nullptr;
@@ -238,6 +278,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) {
                                               &setFixedProgramPtr);
     mock::getCallbacks().set_replace_callback("urProgramRelease",
                                               &releaseFixedProgramPtr);
+    mock::getCallbacks().set_replace_callback("urProgramRetain",
+                                              &customProgramRetain);
 
     const sycl::device Dev = Plt.get_devices()[0];
     sycl::queue Queue{Dev};
@@ -247,8 +289,12 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) {
     auto Mask = PM.getEliminatedKernelArgMask(ProgBefore, Name);
     EXPECT_NE(Mask, nullptr);
     EXPECT_EQ(Mask->at(0), 1);
+    EXPECT_EQ(UsedProgramHandles.size(), 1u);
+    EXPECT_EQ(NativePrograms.count(ProgBefore), 1u);
   }
 
+  EXPECT_EQ(UsedProgramHandles.size(), 0u);
+
   {
     auto Name = sycl::detail::KernelInfo<EAMTestKernel3>::getName();
     sycl::unittest::UrMock<> Mock;
@@ -257,6 +303,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) {
                                               &setFixedProgramPtr);
     mock::getCallbacks().set_replace_callback("urProgramRelease",
                                               &releaseFixedProgramPtr);
+    mock::getCallbacks().set_replace_callback("urProgramRetain",
+                                              &customProgramRetain);
 
     const sycl::device Dev = Plt.get_devices()[0];
     sycl::queue Queue{Dev};
@@ -266,6 +314,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) {
     auto Mask = PM.getEliminatedKernelArgMask(ProgAfter, Name);
     EXPECT_NE(Mask, nullptr);
     EXPECT_EQ(Mask->at(0), 0);
+    EXPECT_EQ(UsedProgramHandles.size(), 1u);
+    EXPECT_EQ(NativePrograms.count(ProgBefore), 1u);
   }
 
   // Verify that the test is behaving correctly and that the pointer is being