[SYCL] Move kernel id into device kernel info struct

Author: sergey-semenov

sycl/source/detail/device_kernel_info.cpp

@@ -12,8 +12,9 @@ namespace sycl {
 inline namespace _V1 {
 namespace detail {
 
-DeviceKernelInfo::DeviceKernelInfo(const CompileTimeKernelInfoTy &Info)
-    : CompileTimeKernelInfoTy(Info) {}
+DeviceKernelInfo::DeviceKernelInfo(const CompileTimeKernelInfoTy &Info,
+                                   std::optional<sycl::kernel_id> KernelID)
+    : CompileTimeKernelInfoTy{Info}, MKernelID{std::move(KernelID)} {}
 
 template <typename OtherTy>
 inline constexpr bool operator==(const CompileTimeKernelInfoTy &LHS,

sycl/source/detail/device_kernel_info.hpp

@@ -13,6 +13,7 @@
 #include <sycl/detail/compile_time_kernel_info.hpp>
 #include <sycl/detail/spinlock.hpp>
 #include <sycl/detail/ur.hpp>
+#include <sycl/kernel_bundle.hpp>
 
 #include <mutex>
 #include <optional>
@@ -84,12 +85,10 @@ struct FastKernelSubcacheT {
 // information that is uniform between different submissions of the same
 // kernel). Pointers to instances of this class are stored in header function
 // templates as a static variable to avoid repeated runtime lookup overhead.
-// TODO Currently this class duplicates information fetched from the program
-// manager. Instead, we should merge all of this information
-// into this structure and get rid of the other KernelName -> * maps.
 class DeviceKernelInfo : public CompileTimeKernelInfoTy {
 public:
-  DeviceKernelInfo(const CompileTimeKernelInfoTy &Info);
+  DeviceKernelInfo(const CompileTimeKernelInfoTy &Info,
+                   std::optional<sycl::kernel_id> KernelID = std::nullopt);
 
   void init(std::string_view KernelName);
   void setCompileTimeInfoIfNeeded(const CompileTimeKernelInfoTy &Info);
@@ -100,6 +99,11 @@ class DeviceKernelInfo : public CompileTimeKernelInfoTy {
     return MImplicitLocalArgPos;
   }
 
+  const sycl::kernel_id &getKernelID() const {
+    assert(MKernelID);
+    return *MKernelID;
+  }
+
   // Implicit local argument position is used only for some backends, so this
   // function allows setting it as more images are added.
   void setImplicitLocalArgPos(int Pos);
@@ -109,6 +113,7 @@ class DeviceKernelInfo : public CompileTimeKernelInfoTy {
 
   FastKernelSubcacheT MFastKernelSubcache;
   std::optional<int> MImplicitLocalArgPos;
+  const std::optional<sycl::kernel_id> MKernelID;
 };
 
 } // namespace detail

sycl/source/detail/program_manager/program_manager.cpp

@@ -1333,11 +1333,12 @@ ProgramManager::getDeviceImage(std::string_view KernelName,
 
   const RTDeviceBinaryImage *Img = nullptr;
   {
-    std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
-    if (auto KernelId = m_KernelName2KernelIDs.find(KernelName);
-        KernelId != m_KernelName2KernelIDs.end()) {
-      Img = getBinImageFromMultiMap(m_KernelIDs2BinImage, KernelId->second,
-                                    ContextImpl, DeviceImpl);
+    std::lock_guard<std::mutex> Guard(m_DeviceKernelInfoMapMutex);
+    if (auto It = m_DeviceKernelInfoMap.find(KernelName);
+        It != m_DeviceKernelInfoMap.end()) {
+      Img = getBinImageFromMultiMap(m_KernelIDs2BinImage,
+                                    It->second.getKernelID(), ContextImpl,
+                                    DeviceImpl);
     }
   }
 
@@ -1721,26 +1722,24 @@ void ProgramManager::addImage(sycl_device_binary RawImg,
 
     auto name = EntriesIt->GetName();
 
-    // Skip creating unique kernel ID if it is an exported device
+    // Skip creating device kernel information 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_ExportedSymbolImages.find(name) != m_ExportedSymbolImages.end())
       continue;
 
-    // ... and create a unique kernel ID for the entry
-    auto It = m_KernelName2KernelIDs.find(name);
-    if (It == m_KernelName2KernelIDs.end()) {
+    auto It = m_DeviceKernelInfoMap.find(std::string_view(name));
+    if (It == m_DeviceKernelInfoMap.end()) {
       sycl::kernel_id KernelID = detail::createSyclObjFromImpl<sycl::kernel_id>(
           std::make_shared<detail::kernel_id_impl>(name));
-
-      It = m_KernelName2KernelIDs.emplace_hint(It, name, KernelID);
+      CompileTimeKernelInfoTy DefaultCompileTimeInfo{std::string_view(name)};
+      It = m_DeviceKernelInfoMap.emplace_hint(
+          It, std::piecewise_construct, std::forward_as_tuple(name),
+          std::forward_as_tuple(DefaultCompileTimeInfo, KernelID));
     }
-    m_KernelIDs2BinImage.insert(std::make_pair(It->second, Img.get()));
-    KernelIDs->push_back(It->second);
-
-    CompileTimeKernelInfoTy DefaultCompileTimeInfo{std::string_view(name)};
-    m_DeviceKernelInfoMap.try_emplace(std::string_view(name),
-                                      DefaultCompileTimeInfo);
+    m_KernelIDs2BinImage.insert(
+        std::make_pair(It->second.getKernelID(), Img.get()));
+    KernelIDs->push_back(It->second.getKernelID());
 
     // Keep track of image to kernel name reference count for cleanup.
     m_KernelNameRefCount[name]++;
@@ -1919,9 +1918,10 @@ void ProgramManager::removeImages(sycl_device_binaries DeviceBinary) {
         continue;
       }
 
-      auto Name2IDIt = m_KernelName2KernelIDs.find(Name);
-      if (Name2IDIt != m_KernelName2KernelIDs.end())
-        removeFromMultimapByVal(m_KernelIDs2BinImage, Name2IDIt->second, Img);
+      auto DKIIt = m_DeviceKernelInfoMap.find(Name);
+      assert(DKIIt != m_DeviceKernelInfoMap.end());
+      removeFromMultimapByVal(m_KernelIDs2BinImage, DKIIt->second.getKernelID(),
+                              Img);
 
       auto RefCountIt = m_KernelNameRefCount.find(Name);
       assert(RefCountIt != m_KernelNameRefCount.end());
@@ -1933,10 +1933,8 @@ void ProgramManager::removeImages(sycl_device_binaries DeviceBinary) {
       if (--RefCount == 0) {
         // TODO aggregate all these maps into a single one since their entries
         // share lifetime.
-        m_DeviceKernelInfoMap.erase(Name);
+        m_DeviceKernelInfoMap.erase(DKIIt);
         m_KernelNameRefCount.erase(RefCountIt);
-        if (Name2IDIt != m_KernelName2KernelIDs.end())
-          m_KernelName2KernelIDs.erase(Name2IDIt);
       }
     }
 
@@ -2058,10 +2056,10 @@ std::vector<kernel_id> ProgramManager::getAllSYCLKernelIDs() {
   std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
 
   std::vector<sycl::kernel_id> AllKernelIDs;
-  AllKernelIDs.reserve(m_KernelName2KernelIDs.size());
-  for (std::pair<std::string_view, kernel_id> KernelID :
-       m_KernelName2KernelIDs) {
-    AllKernelIDs.push_back(KernelID.second);
+  AllKernelIDs.reserve(m_DeviceKernelInfoMap.size());
+  for (const std::pair<const std::string_view, DeviceKernelInfo> &Pair :
+       m_DeviceKernelInfoMap) {
+    AllKernelIDs.push_back(Pair.second.getKernelID());
   }
   return AllKernelIDs;
 }

sycl/source/detail/program_manager/program_manager.hpp

@@ -221,19 +221,20 @@ class ProgramManager {
 
   // The function returns the unique SYCL kernel identifier associated with a
   // kernel name or nullopt if there is no such ID.
-  std::optional<kernel_id> tryGetSYCLKernelID(std::string_view KernelName) {
-    std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
+  std::optional<kernel_id>
+  tryGetSYCLKernelID(std::string_view KernelName) const {
+    std::lock_guard<std::mutex> Guard(m_DeviceKernelInfoMapMutex);
 
-    auto KernelID = m_KernelName2KernelIDs.find(KernelName);
-    if (KernelID == m_KernelName2KernelIDs.end())
+    auto It = m_DeviceKernelInfoMap.find(KernelName);
+    if (It == m_DeviceKernelInfoMap.end())
       return std::nullopt;
 
-    return KernelID->second;
+    return It->second.getKernelID();
   }
 
   // The function returns the unique SYCL kernel identifier associated with a
   // kernel name or throws a sycl exception if there is no such ID.
-  kernel_id getSYCLKernelID(std::string_view KernelName) {
+  kernel_id getSYCLKernelID(std::string_view KernelName) const {
     if (std::optional<kernel_id> MaybeKernelID = tryGetSYCLKernelID(KernelName))
       return *MaybeKernelID;
     throw exception(make_error_code(errc::runtime),
@@ -423,17 +424,10 @@ class ProgramManager {
                                    const device_impl &DeviceImpl);
 
 protected:
-  /// The three maps below are used during kernel resolution. Any kernel is
-  /// identified by its name.
   using RTDeviceBinaryImageUPtr = std::unique_ptr<RTDeviceBinaryImage>;
   using DynRTDeviceBinaryImageUPtr = std::unique_ptr<DynRTDeviceBinaryImage>;
-  /// Maps names of kernels to their unique kernel IDs.
-  /// TODO: Use std::unordered_set with transparent hash and equality functions
-  ///       when C++20 is enabled for the runtime library.
-  /// Access must be guarded by the m_KernelIDsMutex mutex.
-  //
-  std::unordered_map<std::string_view, kernel_id> m_KernelName2KernelIDs;
-
+  /// The two maps below are used during kernel resolution. Any kernel is
+  /// identified by its name, its kernel id is stored in m_DeviceKernelInfoMap.
   // Maps KernelIDs to device binary images. There can be more than one image
   // in case of SPIRV + AOT.
   /// Access must be guarded by the m_KernelIDsMutex mutex.
@@ -525,7 +519,7 @@ class ProgramManager {
   std::unordered_map<std::string_view, DeviceKernelInfo> m_DeviceKernelInfoMap;
 
   // Protects m_DeviceKernelInfoMap.
-  std::mutex m_DeviceKernelInfoMapMutex;
+  mutable std::mutex m_DeviceKernelInfoMapMutex;
 
   // Sanitizer type used in device image
   SanitizerType m_SanitizerFoundInImage;

sycl/source/handler.cpp

@@ -505,8 +505,7 @@ detail::EventImplPtr handler::finalize() {
           !KernelBundleImpPtr->tryGetKernel(impl->getKernelName())) {
         detail::device_impl &Dev = impl->get_device();
         kernel_id KernelID =
-            detail::ProgramManager::getInstance().getSYCLKernelID(
-                impl->getKernelName());
+            impl->MKernelData.getDeviceKernelInfoPtr()->getKernelID();
         bool KernelInserted = KernelBundleImpPtr->add_kernel(
             KernelID, detail::createSyclObjFromImpl<device>(Dev));
         // If kernel was not inserted and the bundle is in input mode we try

sycl/unittests/program_manager/Cleanup.cpp

@@ -18,11 +18,6 @@ class ProgramManagerExposed : public sycl::detail::ProgramManager {
     return m_KernelIDs2BinImage;
   }
 
-  std::unordered_map<std::string_view, sycl::kernel_id> &
-  getKernelName2KernelID() {
-    return m_KernelName2KernelIDs;
-  }
-
   std::unordered_map<const sycl::detail::RTDeviceBinaryImage *,
                      std::shared_ptr<std::vector<sycl::kernel_id>>> &
   getBinImage2KernelId() {
@@ -278,9 +273,6 @@ void checkAllInvolvedContainers(ProgramManagerExposed &PM,
                                 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.getExportedSymbolImages(), ExpectedImgCount,