[SYCL] Delete symbol based info with the last image referencing it (#19659)

Prior to this patch, symbol based info (e.g. kernel id, kernel assert
usage or images containing an exported symbol) was deleted whenever an
image referencing it was removed. This is incorrect since multiple
images can contain the same symbol.

While unlikely to cause any problems now (since those images usually all
get removed with one call to `removeImages` after another), this will
cause issues once kernel name based kernel caches start getting cleaned
up in the same manner.

Author: sergey-semenov

sycl/source/detail/program_manager/program_manager.cpp

@@ -2042,6 +2042,9 @@ void ProgramManager::addImage(sycl_device_binary RawImg,
     }
     m_KernelIDs2BinImage.insert(std::make_pair(It->second, Img.get()));
     KernelIDs->push_back(It->second);
+
+    // Keep track of image to kernel name reference count for cleanup.
+    m_KernelNameRefCount[name]++;
   }
 
   cacheKernelUsesAssertInfo(*Img);
@@ -2115,6 +2118,18 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
     addImage(&(DeviceBinary->DeviceBinaries[I]));
 }
 
+template <typename MultimapT, typename KeyT, typename ValT>
+void removeFromMultimapByVal(MultimapT &Map, const KeyT &Key, const ValT &Val,
+                             bool AssertContains = true) {
+  auto [RangeBegin, RangeEnd] = Map.equal_range(Key);
+  auto It = std::find_if(RangeBegin, RangeEnd,
+                         [&](const auto &Pair) { return Pair.second == Val; });
+  if (!AssertContains && It == RangeEnd)
+    return;
+  assert(It != RangeEnd);
+  Map.erase(It);
+}
+
 void ProgramManager::removeImages(sycl_device_binaries DeviceBinary) {
   if (DeviceBinary->NumDeviceBinaries == 0)
     return;
@@ -2140,44 +2155,68 @@ void ProgramManager::removeImages(sycl_device_binaries DeviceBinary) {
     // Unmap the unique kernel IDs for the offload entries
     for (sycl_offload_entry EntriesIt = EntriesB; EntriesIt != EntriesE;
          EntriesIt = EntriesIt->Increment()) {
-
+      detail::KernelNameStrT Name = EntriesIt->GetName();
       // Drop entry for service kernel
-      if (std::strstr(EntriesIt->GetName(), "__sycl_service_kernel__")) {
-        m_ServiceKernels.erase(EntriesIt->GetName());
+      if (Name.find("__sycl_service_kernel__") != std::string::npos) {
+        removeFromMultimapByVal(m_ServiceKernels, Name, Img);
         continue;
       }
 
       // Exported device functions won't have a kernel ID
-      if (m_ExportedSymbolImages.find(EntriesIt->GetName()) !=
+      if (m_ExportedSymbolImages.find(std::string(Name)) !=
           m_ExportedSymbolImages.end()) {
         continue;
       }
 
-      // remove everything associated with this KernelName
-      m_KernelUsesAssert.erase(EntriesIt->GetName());
-      m_KernelImplicitLocalArgPos.erase(EntriesIt->GetName());
-
-      if (auto It = m_KernelName2KernelIDs.find(EntriesIt->GetName());
-          It != m_KernelName2KernelIDs.end()) {
-        m_KernelIDs2BinImage.erase(It->second);
-        m_KernelName2KernelIDs.erase(It);
+      auto Name2IDIt = m_KernelName2KernelIDs.find(Name);
+      if (Name2IDIt != m_KernelName2KernelIDs.end())
+        removeFromMultimapByVal(m_KernelIDs2BinImage, Name2IDIt->second, Img);
+
+      auto RefCountIt = m_KernelNameRefCount.find(Name);
+      assert(RefCountIt != m_KernelNameRefCount.end());
+      int &RefCount = RefCountIt->second;
+      assert(RefCount > 0);
+
+      // Remove everything associated with this KernelName if this is the last
+      // image referencing it.
+      if (--RefCount == 0) {
+        // TODO aggregate all these maps into a single one since their entries
+        // share lifetime.
+        m_KernelUsesAssert.erase(Name);
+        m_KernelImplicitLocalArgPos.erase(Name);
+        m_KernelNameRefCount.erase(RefCountIt);
+        if (Name2IDIt != m_KernelName2KernelIDs.end())
+          m_KernelName2KernelIDs.erase(Name2IDIt);
       }
     }
 
     // Drop reverse mapping
     m_BinImg2KernelIDs.erase(Img);
 
-    // Unregister exported symbols (needs to happen after the ID unmap loop)
+    // Unregister exported symbol -> Img pair (needs to happen after the ID
+    // unmap loop)
     for (const sycl_device_binary_property &ESProp :
          Img->getExportedSymbols()) {
-      m_ExportedSymbolImages.erase(ESProp->Name);
+      removeFromMultimapByVal(m_ExportedSymbolImages, ESProp->Name, Img,
+                              /*AssertContains*/ false);
     }
 
     for (const sycl_device_binary_property &VFProp :
          Img->getVirtualFunctions()) {
       std::string StrValue = DeviceBinaryProperty(VFProp).asCString();
-      for (const auto &SetName : detail::split_string(StrValue, ','))
-        m_VFSet2BinImage.erase(SetName);
+      // Unregister the image from all referenced virtual function sets.
+      for (const auto &VFSetName : detail::split_string(StrValue, ',')) {
+        auto It = m_VFSet2BinImage.find(VFSetName);
+        assert(It != m_VFSet2BinImage.end());
+        std::set<const RTDeviceBinaryImage *> &ImgSet = It->second;
+        auto ImgIt = ImgSet.find(Img);
+        assert(ImgIt != ImgSet.end());
+        ImgSet.erase(ImgIt);
+        // If no images referencing this virtual function set remain, drop
+        // it from the map.
+        if (ImgSet.empty())
+          m_VFSet2BinImage.erase(It);
+      }
     }
 
     m_DeviceGlobals.eraseEntries(Img);

sycl/source/detail/program_manager/program_manager.hpp

@@ -459,6 +459,12 @@ class ProgramManager {
   /// \ref Sync::getGlobalLock() while holding this mutex.
   std::mutex m_KernelIDsMutex;
 
+  /// Keeps track of binary image to kernel name reference count.
+  /// Used for checking if the last image referencing the kernel name
+  /// is removed in order to trigger cleanup of kernel name based information.
+  /// Access must be guarded by the m_KernelIDsMutex mutex.
+  std::unordered_map<KernelNameStrT, int> m_KernelNameRefCount;
+
   /// Caches all found service kernels to expedite future checks. A SYCL service
   /// kernel is a kernel that has not been defined by the user but is instead
   /// generated by the SYCL runtime. Service kernel name types must be declared

sycl/unittests/program_manager/Cleanup.cpp

@@ -61,6 +61,11 @@ class ProgramManagerExposed : public sycl::detail::ProgramManager {
     return NativePrograms;
   }
 
+  std::unordered_map<sycl::detail::KernelNameStrT, int> &
+  getKernelNameRefCount() {
+    return m_KernelNameRefCount;
+  }
+
   std::unordered_map<const sycl::detail::RTDeviceBinaryImage *,
                      std::unordered_map<sycl::detail::KernelNameStrT,
                                         sycl::detail::KernelArgMask>> &
@@ -132,6 +137,16 @@ std::string generateRefName(const std::string &ImageId,
   return FeatureName + "_" + ImageId;
 }
 
+std::vector<std::string>
+generateRefNames(const std::vector<std::string> &ImageIds,
+                 const std::string &FeatureName) {
+  std::vector<std::string> RefNames;
+  RefNames.reserve(ImageIds.size());
+  for (const std::string &ImageId : ImageIds)
+    RefNames.push_back(generateRefName(ImageId, FeatureName));
+  return RefNames;
+}
+
 sycl::ext::oneapi::experimental::device_global<int> DeviceGlobalA;
 sycl::ext::oneapi::experimental::device_global<int> DeviceGlobalB;
 sycl::ext::oneapi::experimental::device_global<int> DeviceGlobalC;
@@ -143,7 +158,8 @@ using PipeA = sycl::ext::intel::experimental::pipe<PipeIDA, int, 10>;
 using PipeB = sycl::ext::intel::experimental::pipe<PipeIDB, int, 10>;
 using PipeC = sycl::ext::intel::experimental::pipe<PipeIDC, int, 10>;
 
-sycl::unittest::MockDeviceImage generateImage(const std::string &ImageId) {
+sycl::unittest::MockDeviceImage generateImage(const std::string &ImageId,
+                                              bool AddHostPipes = true) {
   sycl::unittest::MockPropertySet PropSet;
 
   std::initializer_list<std::string> KernelNames{
@@ -181,11 +197,11 @@ sycl::unittest::MockDeviceImage generateImage(const std::string &ImageId) {
       std::vector<sycl::unittest::MockProperty>{
           sycl::unittest::makeDeviceGlobalInfo(
               generateRefName(ImageId, "DeviceGlobal"), sizeof(int), 0)});
-
-  PropSet.insert(__SYCL_PROPERTY_SET_SYCL_HOST_PIPES,
-                 std::vector<sycl::unittest::MockProperty>{
-                     sycl::unittest::makeHostPipeInfo(
-                         generateRefName(ImageId, "HostPipe"), sizeof(int))});
+  if (AddHostPipes)
+    PropSet.insert(__SYCL_PROPERTY_SET_SYCL_HOST_PIPES,
+                   std::vector<sycl::unittest::MockProperty>{
+                       sycl::unittest::makeHostPipeInfo(
+                           generateRefName(ImageId, "HostPipe"), sizeof(int))});
   std::vector<unsigned char> Bin{0};
 
   std::vector<sycl::unittest::MockOffloadEntry> Entries =
@@ -229,6 +245,11 @@ static std::array<sycl::unittest::MockDeviceImage, 2> ImagesToKeep = {
 static std::array<sycl::unittest::MockDeviceImage, 1> ImagesToRemove = {
     generateImage("C")};
 
+static std::array<sycl::unittest::MockDeviceImage, 1> ImagesToKeepSameEntries =
+    {generateImage("A", /*AddHostPipe*/ false)};
+static std::array<sycl::unittest::MockDeviceImage, 1>
+    ImagesToRemoveSameEntries = {generateImage("A", /*AddHostPipe*/ false)};
+
 static std::array<sycl::unittest::MockDeviceImage, 2> ImagesToKeepKernelOnly = {
     generateImageKernelOnly("A"), generateImageKernelOnly("B")};
 static std::array<sycl::unittest::MockDeviceImage, 1> ImagesToRemoveKernelOnly =
@@ -251,76 +272,75 @@ void convertAndAddImages(
   PM.addImages(&AllBinaries);
 }
 
-void checkAllInvolvedContainers(ProgramManagerExposed &PM, size_t ExpectedCount,
-                                const std::string &Comment) {
-  EXPECT_EQ(PM.getKernelID2BinImage().size(), ExpectedCount) << Comment;
-  {
-    EXPECT_EQ(PM.getKernelName2KernelID().size(), ExpectedCount) << Comment;
-    EXPECT_TRUE(
-        PM.getKernelName2KernelID().count(generateRefName("A", "Kernel")) > 0)
-        << Comment;
-    EXPECT_TRUE(
-        PM.getKernelName2KernelID().count(generateRefName("B", "Kernel")) > 0)
-        << Comment;
-  }
-  EXPECT_EQ(PM.getBinImage2KernelId().size(), ExpectedCount) << Comment;
-  {
-    EXPECT_EQ(PM.getServiceKernels().size(), ExpectedCount) << Comment;
-    EXPECT_TRUE(PM.getServiceKernels().count(
-                    generateRefName("A", "__sycl_service_kernel__")) > 0)
-        << Comment;
-    EXPECT_TRUE(PM.getServiceKernels().count(
-                    generateRefName("B", "__sycl_service_kernel__")) > 0)
-        << Comment;
-  }
-  {
-    EXPECT_EQ(PM.getExportedSymbolImages().size(), ExpectedCount) << Comment;
-    EXPECT_TRUE(PM.getExportedSymbolImages().count(
-                    generateRefName("A", "Exported")) > 0)
-        << Comment;
-    EXPECT_TRUE(PM.getExportedSymbolImages().count(
-                    generateRefName("B", "Exported")) > 0)
-        << Comment;
-  }
-  EXPECT_EQ(PM.getDeviceImages().size(), ExpectedCount) << Comment;
-  {
-    EXPECT_EQ(PM.getVFSet2BinImage().size(), ExpectedCount) << Comment;
-    EXPECT_TRUE(PM.getVFSet2BinImage().count(generateRefName("A", "VF")) > 0)
-        << Comment;
-    EXPECT_TRUE(PM.getVFSet2BinImage().count(generateRefName("B", "VF")) > 0)
-        << Comment;
+template <typename T>
+void checkContainer(const T &Container, size_t ExpectedCount,
+                    const std::vector<std::string> &ExpectedEntries,
+                    const std::string &Comment) {
+  EXPECT_EQ(Container.size(), ExpectedCount) << Comment;
+  for (const std::string &Entry : ExpectedEntries) {
+    EXPECT_TRUE(Container.count(Entry) > 0) << Comment;
   }
+}
 
-  EXPECT_EQ(PM.getEliminatedKernelArgMask().size(), ExpectedCount) << Comment;
-  {
-    EXPECT_EQ(PM.getKernelUsesAssert().size(), ExpectedCount) << Comment;
-    EXPECT_TRUE(PM.getKernelUsesAssert().count(generateRefName("A", "Kernel")) >
-                0)
-        << Comment;
-    EXPECT_TRUE(PM.getKernelUsesAssert().count(generateRefName("B", "Kernel")) >
-                0)
-        << Comment;
-  }
-  EXPECT_EQ(PM.getKernelImplicitLocalArgPos().size(), ExpectedCount) << Comment;
-
-  {
-    sycl::detail::DeviceGlobalMap &DeviceGlobalMap = PM.getDeviceGlobals();
-    EXPECT_EQ(DeviceGlobalMap.size(), ExpectedCount) << Comment;
-    EXPECT_TRUE(DeviceGlobalMap.count(generateRefName("A", "DeviceGlobal")) > 0)
-        << Comment;
-    EXPECT_TRUE(DeviceGlobalMap.count(generateRefName("B", "DeviceGlobal")) > 0)
-        << Comment;
-    EXPECT_EQ(DeviceGlobalMap.getPointerMap().size(), ExpectedCount) << Comment;
+void checkAllInvolvedContainers(ProgramManagerExposed &PM,
+                                size_t ExpectedImgCount,
+                                size_t ExpectedEntryCount,
+                                const std::vector<std::string> &ImgIds,
+                                const std::string &CommentPostfix,
+                                bool MultipleImgsPerEntryTestCase = false) {
+  EXPECT_EQ(PM.getKernelID2BinImage().size(), ExpectedImgCount)
+      << "KernelID2BinImg " + CommentPostfix;
+  checkContainer(PM.getKernelName2KernelID(), ExpectedEntryCount,
+                 generateRefNames(ImgIds, "Kernel"),
+                 "KernelName2KernelID " + CommentPostfix);
+  EXPECT_EQ(PM.getBinImage2KernelId().size(), ExpectedImgCount)
+      << CommentPostfix;
+  checkContainer(PM.getServiceKernels(), ExpectedImgCount,
+                 generateRefNames(ImgIds, "__sycl_service_kernel__"),
+                 "Service kernels " + CommentPostfix);
+  checkContainer(PM.getExportedSymbolImages(), ExpectedImgCount,
+                 generateRefNames(ImgIds, "Exported"),
+                 "Exported symbol images " + CommentPostfix);
+  EXPECT_EQ(PM.getDeviceImages().size(), ExpectedImgCount)
+      << "Device images " + CommentPostfix;
+
+  checkContainer(PM.getVFSet2BinImage(), ExpectedEntryCount,
+                 generateRefNames(ImgIds, "VF"),
+                 "VFSet2BinImage " + CommentPostfix);
+  checkContainer(PM.getKernelNameRefCount(), ExpectedEntryCount,
+                 generateRefNames(ImgIds, "Kernel"),
+                 "Kernel name reference count " + CommentPostfix);
+  EXPECT_EQ(PM.getEliminatedKernelArgMask().size(), ExpectedImgCount)
+      << "Eliminated kernel arg mask " + CommentPostfix;
+  checkContainer(PM.getKernelUsesAssert(), ExpectedEntryCount,
+                 generateRefNames(ImgIds, "Kernel"),
+                 "KernelUsesAssert " + CommentPostfix);
+  EXPECT_EQ(PM.getKernelImplicitLocalArgPos().size(), ExpectedEntryCount)
+      << "Kernel implicit local arg pos " + CommentPostfix;
+
+  if (!MultipleImgsPerEntryTestCase) {
+    // FIXME expected to fail for now, device globals cleanup seems to be
+    // purging all info for symbols associated with the removed image.
+    checkContainer(PM.getDeviceGlobals(), ExpectedEntryCount,
+                   generateRefNames(ImgIds, "DeviceGlobal"),
+                   "Device globals " + CommentPostfix);
+
+    // The test case with the same entries in multiple images doesn't support
+    // host pipes since those are assumed to be unique.
+    checkContainer(PM.getHostPipes(), ExpectedEntryCount,
+                   generateRefNames(ImgIds, "HostPipe"),
+                   "Host pipes " + CommentPostfix);
+    EXPECT_EQ(PM.getPtrToHostPipe().size(), ExpectedEntryCount)
+        << "Pointer to host pipe " + CommentPostfix;
   }
+}
 
-  {
-    EXPECT_EQ(PM.getHostPipes().size(), ExpectedCount) << Comment;
-    EXPECT_TRUE(PM.getHostPipes().count(generateRefName("A", "HostPipe")) > 0)
-        << Comment;
-    EXPECT_TRUE(PM.getHostPipes().count(generateRefName("B", "HostPipe")) > 0)
-        << Comment;
-  }
-  EXPECT_EQ(PM.getPtrToHostPipe().size(), ExpectedCount) << Comment;
+void checkAllInvolvedContainers(ProgramManagerExposed &PM, size_t ExpectedCount,
+                                const std::vector<std::string> &ImgIds,
+                                const std::string &CommentPostfix,
+                                bool CheckHostPipes = false) {
+  checkAllInvolvedContainers(PM, ExpectedCount, ExpectedCount, ImgIds,
+                             CommentPostfix, CheckHostPipes);
 }
 
 TEST(ImageRemoval, BaseContainers) {
@@ -348,12 +368,37 @@ TEST(ImageRemoval, BaseContainers) {
                             generateRefName("C", "HostPipe").c_str());
 
   checkAllInvolvedContainers(PM, ImagesToRemove.size() + ImagesToKeep.size(),
-                             "Check failed before removal");
+                             {"A", "B", "C"}, "check failed before removal");
+
+  PM.removeImages(&TestBinaries);
+
+  checkAllInvolvedContainers(PM, ImagesToKeep.size(), {"A", "B"},
+                             "check failed after removal");
+}
+
+TEST(ImageRemoval, MultipleImagesPerEntry) {
+  ProgramManagerExposed PM;
+
+  sycl_device_binary_struct NativeImages[ImagesToKeepSameEntries.size()];
+  sycl_device_binaries_struct AllBinaries;
+  convertAndAddImages(PM, ImagesToKeepSameEntries, NativeImages, AllBinaries);
+
+  sycl_device_binary_struct
+      NativeImagesForRemoval[ImagesToRemoveSameEntries.size()];
+  sycl_device_binaries_struct TestBinaries;
+  convertAndAddImages(PM, ImagesToRemoveSameEntries, NativeImagesForRemoval,
+                      TestBinaries);
+
+  checkAllInvolvedContainers(
+      PM, ImagesToRemoveSameEntries.size() + ImagesToKeepSameEntries.size(),
+      /*ExpectedEntryCount*/ 1, {"A"}, "check failed before removal",
+      /*MultipleImgsPerEntryTestCase*/ true);
 
   PM.removeImages(&TestBinaries);
 
-  checkAllInvolvedContainers(PM, ImagesToKeep.size(),
-                             "Check failed after removal");
+  checkAllInvolvedContainers(PM, ImagesToKeepSameEntries.size(), {"A"},
+                             "check failed after removal",
+                             /*MultipleImgsPerEntryTestCase*/ true);
 }
 
 TEST(ImageRemoval, NativePrograms) {