[SYCL] Get rid of device kernel info duplication

With the introduction of DeviceKernelInfo, assert usage and implicit local
argument information is now duplicated in program manager. This patch removes
the duplicate maps and makes it so that device kernel info map is filled out
during image registration, with the compile time information added when it's
available (during the first submission of the kernel).

Author: sergey-semenov

sycl/source/detail/device_kernel_info.cpp

@@ -19,14 +19,10 @@ DeviceKernelInfo::DeviceKernelInfo(const CompileTimeKernelInfoTy &Info)
       Name(Info.Name.data())
 #endif
 {
-  init(Name.data());
-}
-
-void DeviceKernelInfo::init(KernelNameStrRefT KernelName) {
-  auto &PM = detail::ProgramManager::getInstance();
-  MUsesAssert = PM.kernelUsesAssert(KernelName);
-  MImplicitLocalArgPos = PM.kernelImplicitLocalArgPos(KernelName);
 #ifndef __INTEL_PREVIEW_BREAKING_CHANGES
+  // Non-legacy implementation either fills out the data during image
+  // registration after this constructor is called, or uses default values
+  // if this instance of DeviceKernelInfo corresponds to an interop kernel.
   MInitialized.store(true);
 #endif
 }
@@ -36,9 +32,18 @@ void DeviceKernelInfo::initIfEmpty(const CompileTimeKernelInfoTy &Info) {
   if (MInitialized.load())
     return;
 
-  CompileTimeKernelInfoTy::operator=(Info);
-  Name = Info.Name.data();
-  init(Name.data());
+  // If this function is called, then this is a default initialized
+  // device kernel info created from older headers and stored in global handler.
+  // In that case, fetch the proper instance from program manager and copy its
+  // values.
+  auto &PM = detail::ProgramManager::getInstance();
+  DeviceKernelInfo &PMDeviceKernelInfo = PM.getDeviceKernelInfo(KernelNameStrRefT(Info.Name));
+
+  PMDeviceKernelInfo.CompileTimeKernelInfoTy::operator=(Info);
+  PMDeviceKernelInfo.Name = Info.Name.data();
+
+  MUsesAssert = PMDeviceKernelInfo.MUsesAssert;
+  MImplicitLocalArgPos = PMDeviceKernelInfo.MImplicitLocalArgPos;
 }
 #endif
 
@@ -78,18 +83,25 @@ FastKernelSubcacheT &DeviceKernelInfo::getKernelSubcache() {
   assertInitialized();
   return MFastKernelSubcache;
 }
-bool DeviceKernelInfo::usesAssert() {
+bool DeviceKernelInfo::usesAssert() const {
   assertInitialized();
   return MUsesAssert;
 }
-const std::optional<int> &DeviceKernelInfo::getImplicitLocalArgPos() {
+const std::optional<int> &DeviceKernelInfo::getImplicitLocalArgPos() const {
   assertInitialized();
   return MImplicitLocalArgPos;
 }
 
+void DeviceKernelInfo::setUsesAssert() { MUsesAssert = true; }
+
+void DeviceKernelInfo::setImplicitLocalArgPos(int Pos) {
+  assert(!MImplicitLocalArgPos.has_value() || MImplicitLocalArgPos == Pos);
+  MImplicitLocalArgPos = Pos;
+}
+
 bool DeviceKernelInfo::isCompileTimeInfoSet() const { return KernelSize != 0; }
 
-void DeviceKernelInfo::assertInitialized() {
+void DeviceKernelInfo::assertInitialized() const {
 #ifndef __INTEL_PREVIEW_BREAKING_CHANGES
   assert(MInitialized.load() && "Data needs to be initialized before use");
 #endif

sycl/source/detail/device_kernel_info.hpp

@@ -108,11 +108,14 @@ class DeviceKernelInfo : public CompileTimeKernelInfoTy {
   void setCompileTimeInfoIfNeeded(const CompileTimeKernelInfoTy &Info);
 
   FastKernelSubcacheT &getKernelSubcache();
-  bool usesAssert();
-  const std::optional<int> &getImplicitLocalArgPos();
+  bool usesAssert() const;
+  const std::optional<int> &getImplicitLocalArgPos() const;
+
+  void setUsesAssert();
+  void setImplicitLocalArgPos(int Pos);
 
 private:
-  void assertInitialized();
+  void assertInitialized() const;
   bool isCompileTimeInfoSet() const;
 
 #ifndef __INTEL_PREVIEW_BREAKING_CHANGES

sycl/source/detail/get_device_kernel_info.cpp

@@ -22,7 +22,7 @@ KernelNameBasedCacheT *createKernelNameBasedCache() {
 #endif
 
 DeviceKernelInfo &getDeviceKernelInfo(const CompileTimeKernelInfoTy &Info) {
-  return ProgramManager::getInstance().getOrCreateDeviceKernelInfo(Info);
+  return ProgramManager::getInstance().getDeviceKernelInfo(Info);
 }
 
 } // namespace detail

sycl/source/detail/kernel_impl.hpp

@@ -240,10 +240,9 @@ class kernel_impl {
   std::string_view getName() const;
 
   DeviceKernelInfo &getDeviceKernelInfo() {
-    return MIsInterop
-               ? MInteropDeviceKernelInfo
-               : ProgramManager::getInstance().getOrCreateDeviceKernelInfo(
-                     KernelNameStrT(getName()));
+    return MIsInterop ? MInteropDeviceKernelInfo
+                      : ProgramManager::getInstance().getDeviceKernelInfo(
+                            KernelNameStrT(getName()));
   }
 
 private:

sycl/source/detail/program_manager/program_manager.cpp

@@ -1797,8 +1797,11 @@ void ProgramManager::cacheKernelUsesAssertInfo(const RTDeviceBinaryImage &Img) {
   const RTDeviceBinaryImage::PropertyRange &AssertUsedRange =
       Img.getAssertUsed();
   if (AssertUsedRange.isAvailable())
-    for (const auto &Prop : AssertUsedRange)
-      m_KernelUsesAssert.insert(Prop->Name);
+    for (const auto &Prop : AssertUsedRange) {
+      auto It = m_DeviceKernelInfoMap.find(Prop->Name);
+      assert(It != m_DeviceKernelInfoMap.end());
+      It->second.setUsesAssert();
+    }
 }
 
 void ProgramManager::cacheKernelImplicitLocalArg(
@@ -1807,36 +1810,27 @@ void ProgramManager::cacheKernelImplicitLocalArg(
       Img.getImplicitLocalArg();
   if (ImplicitLocalArgRange.isAvailable())
     for (auto Prop : ImplicitLocalArgRange) {
-      m_KernelImplicitLocalArgPos[Prop->Name] =
-          DeviceBinaryProperty(Prop).asUint32();
+      auto It = m_DeviceKernelInfoMap.find(Prop->Name);
+      assert(It != m_DeviceKernelInfoMap.end());
+      It->second.setImplicitLocalArgPos(DeviceBinaryProperty(Prop).asUint32());
     }
 }
 
-std::optional<int>
-ProgramManager::kernelImplicitLocalArgPos(KernelNameStrRefT KernelName) const {
-  auto it = m_KernelImplicitLocalArgPos.find(KernelName);
-  if (it != m_KernelImplicitLocalArgPos.end())
-    return it->second;
-  return {};
-}
-
-DeviceKernelInfo &ProgramManager::getOrCreateDeviceKernelInfo(
-    const CompileTimeKernelInfoTy &Info) {
+DeviceKernelInfo &
+ProgramManager::getDeviceKernelInfo(const CompileTimeKernelInfoTy &Info) {
   std::lock_guard<std::mutex> Guard(m_DeviceKernelInfoMapMutex);
-  auto [Iter, Inserted] =
-      m_DeviceKernelInfoMap.try_emplace(KernelNameStrT{Info.Name.data()}, Info);
-  if (!Inserted)
-    Iter->second.setCompileTimeInfoIfNeeded(Info);
-  return Iter->second;
+  auto It = m_DeviceKernelInfoMap.find(KernelNameStrT{Info.Name.data()});
+  assert(It != m_DeviceKernelInfoMap.end());
+  It->second.setCompileTimeInfoIfNeeded(Info);
+  return It->second;
 }
 
 DeviceKernelInfo &
-ProgramManager::getOrCreateDeviceKernelInfo(KernelNameStrRefT KernelName) {
+ProgramManager::getDeviceKernelInfo(KernelNameStrRefT KernelName) {
   std::lock_guard<std::mutex> Guard(m_DeviceKernelInfoMapMutex);
-  CompileTimeKernelInfoTy DefaultCompileTimeInfo{std::string_view(KernelName)};
-  auto Result =
-      m_DeviceKernelInfoMap.try_emplace(KernelName, DefaultCompileTimeInfo);
-  return Result.first->second;
+  auto It = m_DeviceKernelInfoMap.find(KernelName);
+  assert(It != m_DeviceKernelInfoMap.end());
+  return It->second;
 }
 
 static bool isBfloat16DeviceLibImage(sycl_device_binary RawImg,
@@ -2039,6 +2033,10 @@ void ProgramManager::addImage(sycl_device_binary RawImg,
     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(KernelNameStrT(name),
+                                      DefaultCompileTimeInfo);
+
     // Keep track of image to kernel name reference count for cleanup.
     m_KernelNameRefCount[name]++;
   }
@@ -2232,8 +2230,6 @@ 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_KernelUsesAssert.erase(Name);
-        m_KernelImplicitLocalArgPos.erase(Name);
         m_DeviceKernelInfoMap.erase(Name);
         m_KernelNameRefCount.erase(RefCountIt);
         if (Name2IDIt != m_KernelName2KernelIDs.end())

sycl/source/detail/program_manager/program_manager.hpp

@@ -365,19 +365,13 @@ class ProgramManager {
   ProgramManager();
   ~ProgramManager() = default;
 
-  template <typename NameT>
-  bool kernelUsesAssert(const NameT &KernelName) const {
-    return m_KernelUsesAssert.find(KernelName) != m_KernelUsesAssert.end();
-  }
-
   SanitizerType kernelUsesSanitizer() const { return m_SanitizerFoundInImage; }
 
-  std::optional<int>
-  kernelImplicitLocalArgPos(KernelNameStrRefT KernelName) const;
+  void cacheKernelUsesAssertInfo(const RTDeviceBinaryImage &Img);
+  void cacheKernelImplicitLocalArg(const RTDeviceBinaryImage &Img);
 
-  DeviceKernelInfo &
-  getOrCreateDeviceKernelInfo(const CompileTimeKernelInfoTy &Info);
-  DeviceKernelInfo &getOrCreateDeviceKernelInfo(KernelNameStrRefT KernelName);
+  DeviceKernelInfo &getDeviceKernelInfo(const CompileTimeKernelInfoTy &Info);
+  DeviceKernelInfo &getDeviceKernelInfo(KernelNameStrRefT KernelName);
 
   std::set<const RTDeviceBinaryImage *>
   getRawDeviceImages(const std::vector<kernel_id> &KernelIDs);
@@ -406,12 +400,6 @@ class ProgramManager {
   /// Dumps image to current directory
   void dumpImage(const RTDeviceBinaryImage &Img, uint32_t SequenceID = 0) const;
 
-  /// Add info on kernels using assert into cache
-  void cacheKernelUsesAssertInfo(const RTDeviceBinaryImage &Img);
-
-  /// Add info on kernels using local arg into cache
-  void cacheKernelImplicitLocalArg(const RTDeviceBinaryImage &Img);
-
   std::set<const RTDeviceBinaryImage *>
   collectDependentDeviceImagesForVirtualFunctions(
       const RTDeviceBinaryImage &Img, const device_impl &Dev);
@@ -518,14 +506,6 @@ class ProgramManager {
   bool m_UseSpvFile = false;
   RTDeviceBinaryImageUPtr m_SpvFileImage;
 
-  // std::less<> is a transparent comparator that enabled comparison between
-  // different types without temporary key_type object creation. This includes
-  // standard overloads, such as comparison between std::string and
-  // std::string_view or just char*.
-  using KernelUsesAssertSet = std::set<KernelNameStrT, std::less<>>;
-  KernelUsesAssertSet m_KernelUsesAssert;
-  std::unordered_map<KernelNameStrT, int> m_KernelImplicitLocalArgPos;
-
   // Map for storing device kernel information. Runtime lookup should be avoided
   // by caching the pointers when possible.
   std::unordered_map<KernelNameStrT, DeviceKernelInfo> m_DeviceKernelInfoMap;

sycl/source/handler.cpp

@@ -567,7 +567,7 @@ event handler::finalize() {
       // Fetch the device kernel info pointer if it hasn't been set (e.g.
       // in kernel bundle or free function cases).
       impl->MKernelData.setDeviceKernelInfoPtr(
-          &detail::ProgramManager::getInstance().getOrCreateDeviceKernelInfo(
+          &detail::ProgramManager::getInstance().getDeviceKernelInfo(
               toKernelNameStrT(MKernelName)));
     }
     assert(impl->MKernelData.getKernelName() == MKernelName);
@@ -974,7 +974,7 @@ void handler::extractArgsAndReqs() {
 #ifndef __INTEL_PREVIEW_BREAKING_CHANGES
   if (impl->MKernelData.getDeviceKernelInfoPtr() == nullptr) {
     impl->MKernelData.setDeviceKernelInfoPtr(
-        &detail::ProgramManager::getInstance().getOrCreateDeviceKernelInfo(
+        &detail::ProgramManager::getInstance().getDeviceKernelInfo(
             detail::toKernelNameStrT(MKernel->getName())));
   }
 #endif
@@ -2249,8 +2249,7 @@ void handler::setKernelNameBasedCachePtr(
   HandlerInfo.IsESIMD = impl->MKernelIsESIMD;
   HandlerInfo.HasSpecialCaptures = impl->MKernelHasSpecialCaptures;
   impl->MKernelData.setDeviceKernelInfoPtr(
-      &detail::ProgramManager::getInstance().getOrCreateDeviceKernelInfo(
-          HandlerInfo));
+      &detail::ProgramManager::getInstance().getDeviceKernelInfo(HandlerInfo));
 }
 
 void handler::setKernelInfo(

sycl/unittests/program_manager/Cleanup.cpp

@@ -73,13 +73,6 @@ class ProgramManagerExposed : public sycl::detail::ProgramManager {
     return m_EliminatedKernelArgMasks;
   }
 
-  KernelUsesAssertSet &getKernelUsesAssert() { return m_KernelUsesAssert; }
-
-  std::unordered_map<sycl::detail::KernelNameStrT, int> &
-  getKernelImplicitLocalArgPos() {
-    return m_KernelImplicitLocalArgPos;
-  }
-
   std::unordered_map<std::string,
                      std::unique_ptr<sycl::detail::HostPipeMapEntry>> &
   getHostPipes() {
@@ -311,11 +304,6 @@ void checkAllInvolvedContainers(ProgramManagerExposed &PM,
                  "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
@@ -365,10 +353,6 @@ TEST(ImageRemoval, BaseContainers) {
                             generateRefName("B", "HostPipe").c_str());
   PM.addOrInitHostPipeEntry(PipeC::get_host_ptr(),
                             generateRefName("C", "HostPipe").c_str());
-  std::vector<std::string> KernelNames =
-      generateRefNames({"A", "B", "C"}, "Kernel");
-  for (const std::string &Name : KernelNames)
-    PM.getOrCreateDeviceKernelInfo(Name);
 
   checkAllInvolvedContainers(PM, ImagesToRemove.size() + ImagesToKeep.size(),
                              {"A", "B", "C"}, "check failed before removal");
@@ -392,8 +376,6 @@ TEST(ImageRemoval, MultipleImagesPerEntry) {
   convertAndAddImages(PM, ImagesToRemoveSameEntries, NativeImagesForRemoval,
                       TestBinaries);
 
-  std::string KernelName = generateRefName("A", "Kernel");
-  PM.getOrCreateDeviceKernelInfo(KernelName);
   checkAllInvolvedContainers(
       PM, ImagesToRemoveSameEntries.size() + ImagesToKeepSameEntries.size(),
       /*ExpectedEntryCount*/ 1, {"A"}, "check failed before removal",