diff --git a/sycl/include/CL/sycl/detail/pi.hpp b/sycl/include/CL/sycl/detail/pi.hpp index e1f75f7bf2d6e..13790dbdf5773 100644 --- a/sycl/include/CL/sycl/detail/pi.hpp +++ b/sycl/include/CL/sycl/detail/pi.hpp @@ -369,6 +369,13 @@ class DeviceBinaryImage { return AssertUsed; } const PropertyRange &getProgramMetadata() const { return ProgramMetadata; } + const PropertyRange getExportedSymbols() const { + // We can't have this variable as a class member, since it would break + // the ABI backwards compatibility. + DeviceBinaryImage::PropertyRange ExportedSymbols; + ExportedSymbols.init(Bin, __SYCL_PI_PROPERTY_SET_SYCL_EXPORTED_SYMBOLS); + return ExportedSymbols; + } virtual ~DeviceBinaryImage() {} protected: diff --git a/sycl/source/detail/program_manager/program_manager.cpp b/sycl/source/detail/program_manager/program_manager.cpp index 26171139b6a81..7706318553abb 100644 --- a/sycl/source/detail/program_manager/program_manager.cpp +++ b/sycl/source/detail/program_manager/program_manager.cpp @@ -1059,6 +1059,12 @@ void ProgramManager::addImages(pi_device_binaries DeviceBinary) { KernelSetId KSId = getNextKernelSetId(); { std::lock_guard KernelIDsGuard(m_KernelIDsMutex); + + // Register all exported symbols + auto ExportedSymbols = Img->getExportedSymbols(); + for (const pi_device_binary_property &ExportedSymbol : ExportedSymbols) + m_ExportedSymbols.insert(ExportedSymbol->Name); + for (_pi_offload_entry EntriesIt = EntriesB; EntriesIt != EntriesE; ++EntriesIt) { auto Result = KSIdMap.insert(std::make_pair(EntriesIt->name, KSId)); @@ -1074,6 +1080,13 @@ void ProgramManager::addImages(pi_device_binaries DeviceBinary) { continue; } + // Skip creating unique kernel ID if it is an exported device + // function. Exported device functions appear in the offload entries + // among kernels, but are identifiable by being listed in properties. + if (m_ExportedSymbols.find(EntriesIt->name) != + m_ExportedSymbols.end()) + continue; + // ... and create a unique kernel ID for the entry std::shared_ptr KernelIDImpl = std::make_shared(EntriesIt->name); @@ -1373,9 +1386,11 @@ ProgramManager::getSYCLDeviceImagesWithCompatibleState( auto KernelID = m_KernelIDs.find(EntriesIt->name); if (KernelID == m_KernelIDs.end()) { - // Service kernels do not have kernel IDs - assert(m_ServiceKernels.find(EntriesIt->name) != - m_ServiceKernels.end() && + // Service kernels and exported symbols do not have kernel IDs + assert((m_ServiceKernels.find(EntriesIt->name) != + m_ServiceKernels.end() || + m_ExportedSymbols.find(EntriesIt->name) != + m_ExportedSymbols.end()) && "Kernel ID in device binary missing from cache"); continue; } diff --git a/sycl/source/detail/program_manager/program_manager.hpp b/sycl/source/detail/program_manager/program_manager.hpp index 490118221206d..816ac8260b798 100644 --- a/sycl/source/detail/program_manager/program_manager.hpp +++ b/sycl/source/detail/program_manager/program_manager.hpp @@ -299,6 +299,11 @@ class ProgramManager { /// Access must be guarded by the m_KernelIDsMutex mutex. std::unordered_set m_ServiceKernels; + /// Caches all exported symbols to allow faster lookup when excluding these + // from kernel bundles. + /// Access must be guarded by the m_KernelIDsMutex mutex. + std::unordered_set m_ExportedSymbols; + // Keeps track of pi_program to image correspondence. Needed for: // - knowing which specialization constants are used in the program and // injecting their current values before compiling the SPIR-V; the binary