Revert changes to PM and just use prefix to identify kernels

Signed-off-by: Julian Oppermann <[email protected]>

Author: jopperm

sycl/source/detail/jit_compiler.cpp

@@ -1210,7 +1210,7 @@ std::vector<uint8_t> jit_compiler::encodeReqdWorkGroupSize(
 }
 
 sycl_device_binaries jit_compiler::compileSYCL(
-    const std::string &Id, const std::string &SYCLSource,
+    const std::string &CompilationID, const std::string &SYCLSource,
     const std::vector<std::pair<std::string, std::string>> &IncludePairs,
     const std::vector<std::string> &UserArgs, std::string *LogPtr,
     const std::vector<std::string> &RegisteredKernelNames) {
@@ -1227,7 +1227,7 @@ sycl_device_binaries jit_compiler::compileSYCL(
 
   std::string FinalSource = ss.str();
 
-  std::string SYCLFileName = Id + ".cpp";
+  std::string SYCLFileName = CompilationID + ".cpp";
   ::jit_compiler::InMemoryFile SourceFile{SYCLFileName.c_str(),
                                           FinalSource.c_str()};
 
@@ -1255,7 +1255,7 @@ sycl_device_binaries jit_compiler::compileSYCL(
   }
 
   return createDeviceBinaryImage(Result.getBundleInfo(),
-                                 /*OffloadEntryPrefix=*/Id + '$');
+                                 /*OffloadEntryPrefix=*/CompilationID + '$');
 }
 
 } // namespace detail

sycl/source/detail/jit_compiler.hpp

@@ -46,7 +46,7 @@ class jit_compiler {
                            const std::vector<unsigned char> &SpecConstBlob);
 
   sycl_device_binaries compileSYCL(
-      const std::string &Id, const std::string &SYCLSource,
+      const std::string &CompilationID, const std::string &SYCLSource,
       const std::vector<std::pair<std::string, std::string>> &IncludePairs,
       const std::vector<std::string> &UserArgs, std::string *LogPtr,
       const std::vector<std::string> &RegisteredKernelNames);

sycl/source/detail/kernel_bundle_impl.hpp

@@ -479,34 +479,30 @@ class kernel_bundle_impl {
       // TODO: Support persistent caching.
 
       const std::string &SourceStr = std::get<std::string>(this->Source);
-      auto [Binaries, Id] = syclex::detail::SYCL_JIT_to_SPIRV(
+      auto [Binaries, CompilationID] = syclex::detail::SYCL_JIT_to_SPIRV(
           SourceStr, IncludePairs, BuildOptions, LogPtr, RegisteredKernelNames);
 
       assert(Binaries->NumDeviceBinaries == 1);
 
       auto &PM = detail::ProgramManager::getInstance();
-      std::unordered_set<uintptr_t> ImageIds;
-      PM.addImages(Binaries, &ImageIds);
-      auto DevImgs = PM.getSYCLDeviceImages(
-          MContext, MDevices,
-          [&ImageIds](const detail::DeviceImageImplPtr &DevImgImpl) -> bool {
-            return ImageIds.count(
-                DevImgImpl->get_bin_image_ref()->getImageID());
-          },
-          bundle_state::executable);
-
-      PM.bringSYCLDeviceImagesToState(DevImgs, bundle_state::executable);
+      PM.addImages(Binaries);
 
+      std::vector<kernel_id> KernelIDs;
       std::vector<std::string> KernelNames;
-      std::transform(Binaries->DeviceBinaries->EntriesBegin,
-                     Binaries->DeviceBinaries->EntriesEnd,
-                     std::back_inserter(KernelNames),
-                     [PrefixLen = Id.length() + 1](auto &OffloadEntry) {
-                       // `jit_compiler::compileSYCL` uses `Id + '$'` as name
-                       // prefix; drop that here.
-                       return std::string{OffloadEntry.name + PrefixLen};
-                     });
+      // `jit_compiler::compileSYCL(..)` uses `CompilationID + '$'` as prefix
+      // for offload entry names.
+      std::string Prefix = CompilationID + '$';
+      for (const auto &KernelID : PM.getAllSYCLKernelIDs()) {
+        std::string_view KernelName{KernelID.get_name()};
+        if (KernelName.find(Prefix) == 0) {
+          KernelIDs.push_back(KernelID);
+          KernelName.remove_prefix(Prefix.length());
+          KernelNames.emplace_back(KernelName);
+        }
+      }
 
+      auto DevImgs = PM.getSYCLDeviceImages(MContext, MDevices, KernelIDs,
+                                            bundle_state::executable);
       assert(DevImgs.size() == 1);
       assert(!DevImgs.front().hasDeps());
 

sycl/source/detail/kernel_compiler/kernel_compiler_sycl.cpp

@@ -350,13 +350,13 @@ std::pair<sycl_device_binaries, std::string> SYCL_JIT_to_SPIRV(
     [[maybe_unused]] std::string *LogPtr,
     [[maybe_unused]] const std::vector<std::string> &RegisteredKernelNames) {
 #if SYCL_EXT_JIT_ENABLE
-  static std::atomic_uintptr_t CompilationId;
-  std::string Id = "rtc_" + std::to_string(CompilationId++);
+  static std::atomic_uintptr_t CompilationCounter;
+  std::string CompilationID = "rtc_" + std::to_string(CompilationCounter++);
   sycl_device_binaries Binaries =
       sycl::detail::jit_compiler::get_instance().compileSYCL(
-          Id, SYCLSource, IncludePairs, UserArgs, LogPtr,
+          CompilationID, SYCLSource, IncludePairs, UserArgs, LogPtr,
           RegisteredKernelNames);
-  return std::make_pair(Binaries, std::move(Id));
+  return std::make_pair(Binaries, std::move(CompilationID));
 #else
   throw sycl::exception(sycl::errc::build,
                         "kernel_compiler via sycl-jit is not available");

sycl/source/detail/program_manager/program_manager.cpp

@@ -1802,8 +1802,7 @@ ProgramManager::kernelImplicitLocalArgPos(const std::string &KernelName) const {
   return {};
 }
 
-void ProgramManager::addImages(sycl_device_binaries DeviceBinary,
-                               std::unordered_set<uintptr_t> *ImageIds) {
+void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
   const bool DumpImages = std::getenv("SYCL_DUMP_IMAGES") && !m_UseSpvFile;
   for (int I = 0; I < DeviceBinary->NumDeviceBinaries; I++) {
     sycl_device_binary RawImg = &(DeviceBinary->DeviceBinaries[I]);
@@ -1826,9 +1825,6 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary,
     else
       Img = std::make_unique<RTDeviceBinaryImage>(RawImg);
 
-    if (ImageIds)
-      ImageIds->insert(Img->getImageID());
-
     static uint32_t SequenceID = 0;
 
     // Fill the kernel argument mask map

sycl/source/detail/program_manager/program_manager.hpp

@@ -210,8 +210,7 @@ class ProgramManager {
   ur_program_handle_t getUrProgramFromUrKernel(ur_kernel_handle_t Kernel,
                                                const ContextImplPtr Context);
 
-  void addImages(sycl_device_binaries DeviceImages,
-                 std::unordered_set<uintptr_t> *ImageIds = nullptr);
+  void addImages(sycl_device_binaries DeviceImages);
   void debugPrintBinaryImages() const;
   static std::string getProgramBuildLog(const ur_program_handle_t &Program,
                                         const ContextImplPtr Context);

sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp

@@ -194,6 +194,23 @@ int test_build_and_run() {
   test_1(q, k, 37 + 5);  // ff_cp seeds 37. AddEm will add 5 more.
   test_1(q, k2, 38 + 6); // ff_templated seeds 38. PlusEm adds 6 more.
 
+  // Create and compile new bundle with different header.
+  std::string AddEmHModified = AddEmH;
+  AddEmHModified[AddEmHModified.find('5')] = '7';
+  syclex::include_files incFiles2{"intermediate/AddEm.h", AddEmHModified};
+  incFiles2.add("intermediate/PlusEm.h", PlusEmH);
+  source_kb kbSrc2 = syclex::create_kernel_bundle_from_source(
+      ctx, syclex::source_language::sycl_jit, SYCLSource,
+      syclex::properties{incFiles2});
+
+  exe_kb kbExe3 = syclex::build(kbSrc2);
+  sycl::kernel k3 = kbExe3.ext_oneapi_get_kernel("ff_cp");
+  test_1(q, k3, 37 + 7);
+
+  // Can we still run the original compilation?
+  sycl::kernel k4 = kbExe1.ext_oneapi_get_kernel("ff_cp");
+  test_1(q, k4, 37 + 5);
+
   return 0;
 }