[SYCL][RTC] Experimental use of program manager to build device images

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

Author: jopperm

sycl/source/detail/jit_compiler.cpp

@@ -1120,13 +1120,17 @@ sycl_device_binaries jit_compiler::createPIDeviceBinary(
 }
 
 sycl_device_binaries jit_compiler::createDeviceBinaryImage(
-    const ::jit_compiler::RTCBundleInfo &BundleInfo) {
+    const ::jit_compiler::RTCBundleInfo &BundleInfo,
+    const std::string &OffloadEntryPrefix) {
   DeviceBinaryContainer Binary;
   for (const auto &Symbol : BundleInfo.SymbolTable) {
-    // Create an offload entry for each kernel.
+    // Create an offload entry for each kernel. We prepend a unique prefix to
+    // support reusing the same name across multiple RTC requests. The actual
+    // entrypoints remain unchanged.
     // It seems to be OK to set zero for most of the information here, at least
     // that is the case for compiled SPIR-V binaries.
-    OffloadEntryContainer Entry{Symbol.c_str(), /*Addr=*/nullptr, /*Size=*/0,
+    std::string PrefixedName = OffloadEntryPrefix + Symbol.c_str();
+    OffloadEntryContainer Entry{PrefixedName, /*Addr=*/nullptr, /*Size=*/0,
                                 /*Flags=*/0, /*Reserved=*/0};
     Binary.addOffloadEntry(std::move(Entry));
   }
@@ -1250,7 +1254,8 @@ sycl_device_binaries jit_compiler::compileSYCL(
     throw sycl::exception(sycl::errc::build, Result.getBuildLog());
   }
 
-  return createDeviceBinaryImage(Result.getBundleInfo());
+  return createDeviceBinaryImage(Result.getBundleInfo(),
+                                 /*OffloadEntryPrefix=*/Id + '$');
 }
 
 } // namespace detail

sycl/source/detail/jit_compiler.hpp

@@ -71,7 +71,8 @@ class jit_compiler {
                        ::jit_compiler::BinaryFormat Format);
 
   sycl_device_binaries
-  createDeviceBinaryImage(const ::jit_compiler::RTCBundleInfo &BundleInfo);
+  createDeviceBinaryImage(const ::jit_compiler::RTCBundleInfo &BundleInfo,
+                          const std::string &OffloadEntryPrefix);
 
   std::vector<uint8_t>
   encodeArgUsageMask(const ::jit_compiler::ArgUsageMask &Mask) const;

sycl/source/detail/kernel_bundle_impl.hpp

@@ -474,6 +474,46 @@ class kernel_bundle_impl {
       DeviceVec.push_back(Dev);
     }
 
+    if (Language == syclex::source_language::sycl_jit) {
+      // Experimental: Build device images via the program manager.
+      // TODO: Support persistent caching.
+
+      const std::string &SourceStr = std::get<std::string>(this->Source);
+      auto [Binaries, Id] = 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);
+
+      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};
+                     });
+
+      assert(DevImgs.size() == 1);
+      assert(!DevImgs.front().hasDeps());
+
+      return std::make_shared<kernel_bundle_impl>(
+          MContext, MDevices, DevImgs.front().getMain(), KernelNames, Language);
+    }
+
     ur_program_handle_t UrProgram = nullptr;
     // SourceStrPtr will be null when source is Spir-V bytes.
     const std::string *SourceStrPtr = std::get_if<std::string>(&this->Source);
@@ -514,15 +554,6 @@ class kernel_bundle_impl {
                                                BuildOptions, LogPtr,
                                                RegisteredKernelNames);
         }
-        if (Language == syclex::source_language::sycl_jit) {
-          auto *Binaries = syclex::detail::SYCL_JIT_to_SPIRV(
-              *SourceStrPtr, IncludePairs, BuildOptions, LogPtr,
-              RegisteredKernelNames);
-          assert(Binaries->NumDeviceBinaries == 1 &&
-                 "Device code splitting is not yet supported");
-          return std::vector<uint8_t>(Binaries->DeviceBinaries->BinaryStart,
-                                      Binaries->DeviceBinaries->BinaryEnd);
-        }
         throw sycl::exception(
             make_error_code(errc::invalid),
             "SYCL C++, OpenCL C and SPIR-V are the only supported "

sycl/source/detail/kernel_compiler/kernel_compiler_sycl.cpp

@@ -327,6 +327,7 @@ bool SYCL_Compilation_Available() {
 
 #if SYCL_EXT_JIT_ENABLE
 #include "../jit_compiler.hpp"
+#include <atomic>
 #endif
 
 namespace sycl {
@@ -342,15 +343,20 @@ bool SYCL_JIT_Compilation_Available() {
 #endif
 }
 
-sycl_device_binaries SYCL_JIT_to_SPIRV(
+std::pair<sycl_device_binaries, std::string> SYCL_JIT_to_SPIRV(
     [[maybe_unused]] const std::string &SYCLSource,
     [[maybe_unused]] include_pairs_t IncludePairs,
     [[maybe_unused]] const std::vector<std::string> &UserArgs,
     [[maybe_unused]] std::string *LogPtr,
     [[maybe_unused]] const std::vector<std::string> &RegisteredKernelNames) {
 #if SYCL_EXT_JIT_ENABLE
-  return sycl::detail::jit_compiler::get_instance().compileSYCL(
-      "rtc", SYCLSource, IncludePairs, UserArgs, LogPtr, RegisteredKernelNames);
+  static std::atomic_uintptr_t CompilationId;
+  std::string Id = "rtc_" + std::to_string(CompilationId++);
+  sycl_device_binaries Binaries =
+      sycl::detail::jit_compiler::get_instance().compileSYCL(
+          Id, SYCLSource, IncludePairs, UserArgs, LogPtr,
+          RegisteredKernelNames);
+  return std::make_pair(Binaries, std::move(Id));
 #else
   throw sycl::exception(sycl::errc::build,
                         "kernel_compiler via sycl-jit is not available");

sycl/source/detail/kernel_compiler/kernel_compiler_sycl.hpp

@@ -35,7 +35,7 @@ bool SYCL_Compilation_Available();
 
 std::string userArgsAsString(const std::vector<std::string> &UserArguments);
 
-sycl_device_binaries
+std::pair<sycl_device_binaries, std::string>
 SYCL_JIT_to_SPIRV(const std::string &Source, include_pairs_t IncludePairs,
                   const std::vector<std::string> &UserArgs, std::string *LogPtr,
                   const std::vector<std::string> &RegisteredKernelNames);

sycl/source/detail/program_manager/program_manager.cpp

@@ -1802,7 +1802,8 @@ ProgramManager::kernelImplicitLocalArgPos(const std::string &KernelName) const {
   return {};
 }
 
-void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
+void ProgramManager::addImages(sycl_device_binaries DeviceBinary,
+                               std::unordered_set<uintptr_t> *ImageIds) {
   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]);
@@ -1825,6 +1826,9 @@ 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,7 +210,8 @@ class ProgramManager {
   ur_program_handle_t getUrProgramFromUrKernel(ur_kernel_handle_t Kernel,
                                                const ContextImplPtr Context);
 
-  void addImages(sycl_device_binaries DeviceImages);
+  void addImages(sycl_device_binaries DeviceImages,
+                 std::unordered_set<uintptr_t> *ImageIds = nullptr);
   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

@@ -13,26 +13,6 @@
 // RUN: %{run} %t.out 1
 // RUN: %{l0_leak_check} %{run} %t.out 1
 
-// -- Test again, with caching.
-
-// DEFINE: %{cache_vars} = %{l0_leak_check} env SYCL_CACHE_PERSISTENT=1 SYCL_CACHE_TRACE=5 SYCL_CACHE_DIR=%t/cache_dir
-// RUN: %if run-mode %{ rm -rf %t/cache_dir %}
-// RUN: %{cache_vars} %{run-unfiltered-devices} %t.out 2>&1 |  FileCheck %s --check-prefixes=CHECK-WRITTEN-TO-CACHE
-// RUN: %{cache_vars} %{run-unfiltered-devices} %t.out 2>&1 |  FileCheck %s --check-prefixes=CHECK-READ-FROM-CACHE
-
-// -- Add leak check.
-// RUN: %if run-mode %{ rm -rf %t/cache_dir %}
-// RUN: %{l0_leak_check} %{cache_vars} %{run-unfiltered-devices} %t.out 2>&1 |  FileCheck %s --check-prefixes=CHECK-WRITTEN-TO-CACHE
-// RUN: %{l0_leak_check} %{cache_vars} %{run-unfiltered-devices} %t.out 2>&1 |  FileCheck %s --check-prefixes=CHECK-READ-FROM-CACHE
-
-// CHECK-WRITTEN-TO-CACHE: [Persistent Cache]: enabled
-// CHECK-WRITTEN-TO-CACHE-NOT: [kernel_compiler Persistent Cache]: using cached binary
-// CHECK-WRITTEN-TO-CACHE: [kernel_compiler Persistent Cache]: binary has been cached
-
-// CHECK-READ-FROM-CACHE: [Persistent Cache]: enabled
-// CHECK-READ-FROM-CACHE-NOT: [kernel_compiler Persistent Cache]: binary has been cached
-// CHECK-READ-FROM-CACHE: [kernel_compiler Persistent Cache]: using cached binary
-
 #include <sycl/detail/core.hpp>
 #include <sycl/kernel_bundle.hpp>
 #include <sycl/usm.hpp>