diff --git a/sycl/source/detail/program_manager/program_manager.cpp b/sycl/source/detail/program_manager/program_manager.cpp index a7c6d263079e8..f7e2f8477b583 100644 --- a/sycl/source/detail/program_manager/program_manager.cpp +++ b/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 +void removeFromMultimap(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__")) { + removeFromMultimap(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()); + // Remove everything associated with this KernelName if this is the last + // image referencing it, otherwise remove just the ID -> Img mapping. + auto RefCountIt = m_KernelNameRefCount.find(Name); + assert(RefCountIt != m_KernelNameRefCount.end()); + int &RefCount = RefCountIt->second; + assert(RefCount > 0); + --RefCount; - if (auto It = m_KernelName2KernelIDs.find(EntriesIt->GetName()); + if (auto It = m_KernelName2KernelIDs.find(Name); It != m_KernelName2KernelIDs.end()) { - m_KernelIDs2BinImage.erase(It->second); - m_KernelName2KernelIDs.erase(It); + if (RefCount == 0) { + m_KernelIDs2BinImage.erase(It->second); + m_KernelName2KernelIDs.erase(It); + } else { + removeFromMultimap(m_KernelIDs2BinImage, It->second, Img); + } + } + + if (RefCount == 0) { + m_KernelUsesAssert.erase(Name); + m_KernelImplicitLocalArgPos.erase(Name); + m_KernelNameRefCount.erase(RefCountIt); } } // 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); + removeFromMultimap(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); + for (const auto &SetName : detail::split_string(StrValue, ',')) { + auto It = m_VFSet2BinImage.find(SetName); + assert(It != m_VFSet2BinImage.end()); + auto &ImgSet = It->second; + auto ImgIt = ImgSet.find(Img); + assert(ImgIt != ImgSet.end()); + ImgSet.erase(ImgIt); + if (ImgSet.empty()) + m_VFSet2BinImage.erase(It); + } } m_DeviceGlobals.eraseEntries(Img); diff --git a/sycl/source/detail/program_manager/program_manager.hpp b/sycl/source/detail/program_manager/program_manager.hpp index 56544ec3b9411..eaea458f95e84 100644 --- a/sycl/source/detail/program_manager/program_manager.hpp +++ b/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 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 diff --git a/sycl/unittests/program_manager/Cleanup.cpp b/sycl/unittests/program_manager/Cleanup.cpp index 47be8b7d05645..4a39f06ee5250 100644 --- a/sycl/unittests/program_manager/Cleanup.cpp +++ b/sycl/unittests/program_manager/Cleanup.cpp @@ -61,6 +61,11 @@ class ProgramManagerExposed : public sycl::detail::ProgramManager { return NativePrograms; } + std::unordered_map & + getKernelNameRefCount() { + return m_KernelNameRefCount; + } + std::unordered_map> & @@ -132,6 +137,16 @@ std::string generateRefName(const std::string &ImageId, return FeatureName + "_" + ImageId; } +std::vector +generateRefNames(const std::vector &ImageIds, + const std::string &FeatureName) { + std::vector 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 DeviceGlobalA; sycl::ext::oneapi::experimental::device_global DeviceGlobalB; sycl::ext::oneapi::experimental::device_global DeviceGlobalC; @@ -143,7 +158,8 @@ using PipeA = sycl::ext::intel::experimental::pipe; using PipeB = sycl::ext::intel::experimental::pipe; using PipeC = sycl::ext::intel::experimental::pipe; -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 KernelNames{ @@ -181,11 +197,11 @@ sycl::unittest::MockDeviceImage generateImage(const std::string &ImageId) { std::vector{ sycl::unittest::makeDeviceGlobalInfo( generateRefName(ImageId, "DeviceGlobal"), sizeof(int), 0)}); - - PropSet.insert(__SYCL_PROPERTY_SET_SYCL_HOST_PIPES, - std::vector{ - sycl::unittest::makeHostPipeInfo( - generateRefName(ImageId, "HostPipe"), sizeof(int))}); + if (AddHostPipes) + PropSet.insert(__SYCL_PROPERTY_SET_SYCL_HOST_PIPES, + std::vector{ + sycl::unittest::makeHostPipeInfo( + generateRefName(ImageId, "HostPipe"), sizeof(int))}); std::vector Bin{0}; std::vector Entries = @@ -229,6 +245,11 @@ static std::array ImagesToKeep = { static std::array ImagesToRemove = { generateImage("C")}; +static std::array ImagesToKeepSameEntries = + {generateImage("A", /*AddHostPipe*/ false)}; +static std::array + ImagesToRemoveSameEntries = {generateImage("A", /*AddHostPipe*/ false)}; + static std::array ImagesToKeepKernelOnly = { generateImageKernelOnly("A"), generateImageKernelOnly("B")}; static std::array 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 +void checkContainer(const T &Container, size_t ExpectedCount, + const std::vector &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 &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 &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) {