[SYCL] Delete symbol based info with the last image referencing it

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,67 @@ 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()) {
-
+      const char *Name = EntriesIt->GetName();
       // Drop entry for service kernel
-      if (std::strstr(EntriesIt->GetName(), "__sycl_service_kernel__")) {
-        m_ServiceKernels.erase(EntriesIt->GetName());
+      if (std::strstr(Name, "__sycl_service_kernel__")) {
+        removeFromMultimapByVal(m_ServiceKernels, Name, Img);
         continue;
       }
 
       // Exported device functions won't have a kernel ID
-      if (m_ExportedSymbolImages.find(EntriesIt->GetName()) !=
-          m_ExportedSymbolImages.end()) {
+      if (m_ExportedSymbolImages.find(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) {