[SYCL] Cache 'kernel uses assert' info

sycl/source/detail/program_manager/program_manager.cpp

@@ -999,6 +999,23 @@ createKernelArgMask(const pi::ByteArray &Bytes) {
   return Result;
 }
 
+void ProgramManager::cacheKernelUsesAssertInfo(OSModuleHandle M,
+                                               RTDeviceBinaryImage &Img) {
+  const pi::DeviceBinaryImage::PropertyRange &AssertUsedRange =
+      Img.getAssertUsed();
+  if (AssertUsedRange.isAvailable())
+    for (const auto &Prop : AssertUsedRange) {
+      KernelNameWithOSModule Key{Prop->Name, M};
+      m_KernelUsesAssert.insert(Key);
+    }
+}
+
+bool ProgramManager::kernelUsesAssert(OSModuleHandle M,
+                                      const std::string &KernelName) const {
+  KernelNameWithOSModule Key{KernelName, M};
+  return m_KernelUsesAssert.find(Key) != m_KernelUsesAssert.end();
+}
+
 void ProgramManager::addImages(pi_device_binaries DeviceBinary) {
   std::lock_guard<std::mutex> Guard(Sync::getGlobalLock());
 
@@ -1032,6 +1049,9 @@ void ProgramManager::addImages(pi_device_binaries DeviceBinary) {
                  "Kernel sets are not disjoint");
         auto &Imgs = m_DeviceImages[KSIdIt->second];
         assert(Imgs && "Device image vector should have been already created");
+
+        cacheKernelUsesAssertInfo(M, *Img);
+
         Imgs->push_back(std::move(Img));
         continue;
       }
@@ -1054,6 +1074,9 @@ void ProgramManager::addImages(pi_device_binaries DeviceBinary) {
         }
       }
       m_DeviceImages[KSId].reset(new std::vector<RTDeviceBinaryImageUPtr>());
+
+      cacheKernelUsesAssertInfo(M, *Img);
+
       m_DeviceImages[KSId]->push_back(std::move(Img));
       continue;
     }
@@ -1066,6 +1089,9 @@ void ProgramManager::addImages(pi_device_binaries DeviceBinary) {
     auto &Imgs = m_DeviceImages[KSId];
     if (!Imgs)
       Imgs.reset(new std::vector<RTDeviceBinaryImageUPtr>());
+
+    cacheKernelUsesAssertInfo(M, *Img);
+
     Imgs->push_back(std::move(Img));
   }
 }

sycl/source/detail/program_manager/program_manager.hpp

@@ -21,6 +21,7 @@
 #include <cstdint>
 #include <map>
 #include <memory>
+#include <set>
 #include <unordered_map>
 #include <vector>
 
@@ -212,6 +213,8 @@ class ProgramManager {
   ProgramManager();
   ~ProgramManager() = default;
 
+  bool kernelUsesAssert(OSModuleHandle M, const std::string &KernelName) const;
+
 private:
   ProgramManager(ProgramManager const &) = delete;
   ProgramManager &operator=(ProgramManager const &) = delete;
@@ -237,6 +240,9 @@ class ProgramManager {
   /// Dumps image to current directory
   void dumpImage(const RTDeviceBinaryImage &Img, KernelSetId KSId) const;
 
+  /// Add info on kernels using assert into cache
+  void cacheKernelUsesAssertInfo(OSModuleHandle M, RTDeviceBinaryImage &Img);
+
   /// The three maps below are used during kernel resolution. Any kernel is
   /// identified by its name and the OS module it's coming from, allowing
   /// kernels with identical names in different OS modules. The following
@@ -310,6 +316,9 @@ class ProgramManager {
 
   /// True iff a SPIR-V file has been specified with an environment variable
   bool m_UseSpvFile = false;
+
+  using KernelNameWithOSModule = std::pair<std::string, OSModuleHandle>;
+  std::set<KernelNameWithOSModule> m_KernelUsesAssert;
 };
 } // namespace detail
 } // namespace sycl

sycl/source/detail/queue_impl.cpp

@@ -351,20 +351,6 @@ pi_native_handle queue_impl::getNative() const {
   return Handle;
 }
 
-bool queue_impl::kernelUsesAssert(const std::string &KernelName,
-                                  OSModuleHandle Handle) const {
-  RTDeviceBinaryImage &BinImg = ProgramManager::getInstance().getDeviceImage(
-      Handle, KernelName, get_context(), get_device());
-
-  const pi::DeviceBinaryImage::PropertyRange &AssertUsedRange =
-      BinImg.getAssertUsed();
-  if (AssertUsedRange.isAvailable())
-    for (const auto &Prop : AssertUsedRange)
-      if (Prop->Name == KernelName)
-        return true;
-
-  return false;
-}
 } // namespace detail
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)

sycl/source/detail/queue_impl.hpp

@@ -402,9 +402,6 @@ class queue_impl {
   /// \return a native handle.
   pi_native_handle getNative() const;
 
-  bool kernelUsesAssert(const std::string &KernelName,
-                        OSModuleHandle Handle) const;
-
   buffer<AssertHappened, 1> &getAssertHappenedBuffer() {
     return MAssertHappenedBuffer;
   }
@@ -438,10 +435,10 @@ class queue_impl {
       bool IsKernel = Handler.getType() == CG::Kernel;
       bool KernelUsesAssert = false;
       if (IsKernel)
-        KernelUsesAssert = Handler.MKernel
-                               ? true
-                               : kernelUsesAssert(Handler.MKernelName,
-                                                  Handler.MOSModuleHandle);
+        KernelUsesAssert =
+            Handler.MKernel ? true
+                            : ProgramManager::getInstance().kernelUsesAssert(
+                                  Handler.MOSModuleHandle, Handler.MKernelName);
 
       Event = Handler.finalize();