[SYCL][RTC] Define custom destructor for kernel_bundle_impl (#16702)

Adds a custom destructor to `kernel_bundle_impl` to properly clean-up
runtime information and device binaries for bundles that are
runtime-compiled from source:

- **Removal of kernels from program manager**
The new `ProgramManager::removeImages` method takes `addImages` as a
blueprint and reverses the effects of registering the given device
binaries. Currently, only a subset of the data structures in the program
manager is handled, and `assert`s are in place for the remaining
members. Device globals will be addressed after #16565 lands. AFAICT I
could clean-up most of the other members mechanically as well, but
decided against that because I can't test these features (such as
virtual functions) right now due to lack of SYCL-RTC support.

- **Free device binaries in JIT library**
Store the UR-specific data structure that backs a device binary in a map
of unique ptrs instead of a vector to make it possible to free it again
without invalidating other compilation results. Also introduce a mutex
to protect concurrent access to this new map.

- **Free raw SPIR-V binaries in JIT context**
Again replace a vector with a map of unique ptr to make it possible to
free `KernelBinary` objects in the `JITContext`. `KernelBinary` owns the
actual SPIR-V binaries; all other data structures mentioned earlier only
store the pointer return by `KernelBinary::address()`. A new function
`destroyBinary` is added to the `sycl-jit` interface.

My understanding is that `~device_image_impl` and `~kernel_impl` already
take care of releasing the underlying UR resources, hence no changes
required there.

---------

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

Author: jopperm

sycl-jit/jit-compiler/include/KernelFusion.h

@@ -104,6 +104,8 @@ KF_EXPORT_SYMBOL RTCResult compileSYCL(InMemoryFile SourceFile,
                                        View<InMemoryFile> IncludeFiles,
                                        View<const char *> UserArgs);
 
+KF_EXPORT_SYMBOL void destroyBinary(BinaryAddress Address);
+
 /// Clear all previously set options.
 KF_EXPORT_SYMBOL void resetJITConfiguration();
 

sycl-jit/jit-compiler/ld-version-script.txt

@@ -4,6 +4,7 @@
     fuseKernels;
     materializeSpecConstants;
     compileSYCL;
+    destroyBinary;
     resetJITConfiguration;
     addToJITConfiguration;
 

sycl-jit/jit-compiler/lib/KernelFusion.cpp

@@ -317,6 +317,10 @@ compileSYCL(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
   return RTCResult{std::move(BundleInfo), BuildLog.c_str()};
 }
 
+extern "C" KF_EXPORT_SYMBOL void destroyBinary(BinaryAddress Address) {
+  JITContext::getInstance().destroyKernelBinary(Address);
+}
+
 extern "C" KF_EXPORT_SYMBOL void resetJITConfiguration() {
   ConfigHelper::reset();
 }

sycl-jit/jit-compiler/lib/fusion/JITContext.h

@@ -70,7 +70,15 @@ class JITContext {
 
   template <typename... Ts> KernelBinary &emplaceKernelBinary(Ts &&...Args) {
     WriteLockT WriteLock{BinariesMutex};
-    return Binaries.emplace_back(std::forward<Ts>(Args)...);
+    auto KBUPtr = std::make_unique<KernelBinary>(std::forward<Ts>(Args)...);
+    KernelBinary &KB = *KBUPtr;
+    Binaries.emplace(KB.address(), std::move(KBUPtr));
+    return KB;
+  }
+
+  void destroyKernelBinary(BinaryAddress Addr) {
+    WriteLockT WriteLock{BinariesMutex};
+    Binaries.erase(Addr);
   }
 
   std::optional<SYCLKernelInfo> getCacheEntry(CacheKeyT &Identifier) const;
@@ -96,7 +104,7 @@ class JITContext {
 
   MutexT BinariesMutex;
 
-  std::vector<KernelBinary> Binaries;
+  std::unordered_map<BinaryAddress, std::unique_ptr<KernelBinary>> Binaries;
 
   mutable MutexT CacheMutex;
 

sycl/source/detail/jit_compiler.cpp

@@ -98,6 +98,16 @@ jit_compiler::jit_compiler()
           "Cannot resolve JIT library function entry point");
       return false;
     }
+
+    this->DestroyBinaryHandle = reinterpret_cast<DestroyBinaryFuncT>(
+        sycl::detail::ur::getOsLibraryFuncAddress(LibraryPtr.get(),
+                                                  "destroyBinary"));
+    if (!this->DestroyBinaryHandle) {
+      printPerformanceWarning(
+          "Cannot resolve JIT library function entry point");
+      return false;
+    }
+
     LibraryHandle = std::move(LibraryPtr);
     return true;
   };
@@ -1130,10 +1140,10 @@ sycl_device_binaries jit_compiler::createPIDeviceBinary(
   return JITDeviceBinaries.back().getPIDeviceStruct();
 }
 
-sycl_device_binaries jit_compiler::createDeviceBinaryImage(
+sycl_device_binaries jit_compiler::createDeviceBinaries(
     const ::jit_compiler::RTCBundleInfo &BundleInfo,
     const std::string &OffloadEntryPrefix) {
-  DeviceBinariesCollection Collection;
+  auto Collection = std::make_unique<DeviceBinariesCollection>();
 
   for (const auto &DevImgInfo : BundleInfo) {
     DeviceBinaryContainer Binary;
@@ -1164,17 +1174,28 @@ sycl_device_binaries jit_compiler::createDeviceBinaryImage(
       Binary.addProperty(std::move(PropSet));
     }
 
-    Collection.addDeviceBinary(std::move(Binary),
-                               DevImgInfo.BinaryInfo.BinaryStart,
-                               DevImgInfo.BinaryInfo.BinarySize,
-                               (DevImgInfo.BinaryInfo.AddressBits == 64)
-                                   ? __SYCL_DEVICE_BINARY_TARGET_SPIRV64
-                                   : __SYCL_DEVICE_BINARY_TARGET_SPIRV32,
-                               SYCL_DEVICE_BINARY_TYPE_SPIRV);
+    Collection->addDeviceBinary(std::move(Binary),
+                                DevImgInfo.BinaryInfo.BinaryStart,
+                                DevImgInfo.BinaryInfo.BinarySize,
+                                (DevImgInfo.BinaryInfo.AddressBits == 64)
+                                    ? __SYCL_DEVICE_BINARY_TARGET_SPIRV64
+                                    : __SYCL_DEVICE_BINARY_TARGET_SPIRV32,
+                                SYCL_DEVICE_BINARY_TYPE_SPIRV);
   }
 
-  JITDeviceBinaries.push_back(std::move(Collection));
-  return JITDeviceBinaries.back().getPIDeviceStruct();
+  sycl_device_binaries Binaries = Collection->getPIDeviceStruct();
+
+  std::lock_guard<std::mutex> Guard{RTCDeviceBinariesMutex};
+  RTCDeviceBinaries.emplace(Binaries, std::move(Collection));
+  return Binaries;
+}
+
+void jit_compiler::destroyDeviceBinaries(sycl_device_binaries Binaries) {
+  std::lock_guard<std::mutex> Guard{RTCDeviceBinariesMutex};
+  for (uint16_t i = 0; i < Binaries->NumDeviceBinaries; ++i) {
+    DestroyBinaryHandle(Binaries->DeviceBinaries[i].BinaryStart);
+  }
+  RTCDeviceBinaries.erase(Binaries);
 }
 
 std::vector<uint8_t> jit_compiler::encodeArgUsageMask(
@@ -1270,8 +1291,8 @@ sycl_device_binaries jit_compiler::compileSYCL(
     throw sycl::exception(sycl::errc::build, Result.getBuildLog());
   }
 
-  return createDeviceBinaryImage(Result.getBundleInfo(),
-                                 /*OffloadEntryPrefix=*/CompilationID + '$');
+  return createDeviceBinaries(Result.getBundleInfo(),
+                              /*OffloadEntryPrefix=*/CompilationID + '$');
 }
 
 } // namespace detail

sycl/source/detail/jit_compiler.hpp

@@ -17,6 +17,8 @@
 #endif // SYCL_EXT_JIT_ENABLE
 
 #include <functional>
+#include <memory>
+#include <mutex>
 #include <unordered_map>
 
 namespace jit_compiler {
@@ -53,6 +55,8 @@ class jit_compiler {
       const std::vector<std::string> &UserArgs, std::string *LogPtr,
       const std::vector<std::string> &RegisteredKernelNames);
 
+  void destroyDeviceBinaries(sycl_device_binaries Binaries);
+
   bool isAvailable() { return Available; }
 
   static jit_compiler &get_instance() {
@@ -73,8 +77,8 @@ class jit_compiler {
                        ::jit_compiler::BinaryFormat Format);
 
   sycl_device_binaries
-  createDeviceBinaryImage(const ::jit_compiler::RTCBundleInfo &BundleInfo,
-                          const std::string &OffloadEntryPrefix);
+  createDeviceBinaries(const ::jit_compiler::RTCBundleInfo &BundleInfo,
+                       const std::string &OffloadEntryPrefix);
 
   std::vector<uint8_t>
   encodeArgUsageMask(const ::jit_compiler::ArgUsageMask &Mask) const;
@@ -88,17 +92,27 @@ class jit_compiler {
   // Manages the lifetime of the UR structs for device binaries.
   std::vector<DeviceBinariesCollection> JITDeviceBinaries;
 
+  // Manages the lifetime of the UR structs for device binaries for SYCL-RTC.
+  std::unordered_map<sycl_device_binaries,
+                     std::unique_ptr<DeviceBinariesCollection>>
+      RTCDeviceBinaries;
+
+  // Protects access to map above.
+  std::mutex RTCDeviceBinariesMutex;
+
 #if SYCL_EXT_JIT_ENABLE
   // Handles to the entry points of the lazily loaded JIT library.
   using FuseKernelsFuncT = decltype(::jit_compiler::fuseKernels) *;
   using MaterializeSpecConstFuncT =
       decltype(::jit_compiler::materializeSpecConstants) *;
   using CompileSYCLFuncT = decltype(::jit_compiler::compileSYCL) *;
+  using DestroyBinaryFuncT = decltype(::jit_compiler::destroyBinary) *;
   using ResetConfigFuncT = decltype(::jit_compiler::resetJITConfiguration) *;
   using AddToConfigFuncT = decltype(::jit_compiler::addToJITConfiguration) *;
   FuseKernelsFuncT FuseKernelsHandle = nullptr;
   MaterializeSpecConstFuncT MaterializeSpecConstHandle = nullptr;
   CompileSYCLFuncT CompileSYCLHandle = nullptr;
+  DestroyBinaryFuncT DestroyBinaryHandle = nullptr;
   ResetConfigFuncT ResetConfigHandle = nullptr;
   AddToConfigFuncT AddToConfigHandle = nullptr;
   static std::function<void(void *)> CustomDeleterForLibHandle;

sycl/source/detail/kernel_bundle_impl.hpp

@@ -380,7 +380,8 @@ class kernel_bundle_impl {
   // program manager integration, only for sycl_jit language
   kernel_bundle_impl(context Ctx, std::vector<device> Devs,
                      const std::vector<kernel_id> &KernelIDs,
-                     std::vector<std::string> KNames, std::string Pfx,
+                     std::vector<std::string> KNames,
+                     sycl_device_binaries Binaries, std::string Pfx,
                      syclex::source_language Lang)
       : kernel_bundle_impl(std::move(Ctx), std::move(Devs), KernelIDs,
                            bundle_state::executable) {
@@ -392,6 +393,7 @@ class kernel_bundle_impl {
     // from the (unprefixed) kernel name.
     MIsInterop = true;
     MKernelNames = std::move(KNames);
+    MDeviceBinaries = Binaries;
     MPrefix = std::move(Pfx);
     MLanguage = Lang;
   }
@@ -520,8 +522,9 @@ class kernel_bundle_impl {
         }
       }
 
-      return std::make_shared<kernel_bundle_impl>(
-          MContext, MDevices, KernelIDs, KernelNames, Prefix, MLanguage);
+      return std::make_shared<kernel_bundle_impl>(MContext, MDevices, KernelIDs,
+                                                  KernelNames, Binaries, Prefix,
+                                                  MLanguage);
     }
 
     ur_program_handle_t UrProgram = nullptr;
@@ -928,6 +931,17 @@ class kernel_bundle_impl {
     return true;
   }
 
+  ~kernel_bundle_impl() {
+    try {
+      if (MDeviceBinaries) {
+        ProgramManager::getInstance().removeImages(MDeviceBinaries);
+        syclex::detail::SYCL_JIT_destroy(MDeviceBinaries);
+      }
+    } catch (std::exception &e) {
+      __SYCL_REPORT_EXCEPTION_TO_STREAM("exception in ~kernel_bundle_impl", e);
+    }
+  }
+
 private:
   void fillUniqueDeviceImages() {
     assert(MUniqueDeviceImages.empty());
@@ -959,6 +973,7 @@ class kernel_bundle_impl {
   const std::variant<std::string, std::vector<std::byte>> MSource;
   // only kernel_bundles created from source have KernelNames member.
   std::vector<std::string> MKernelNames;
+  sycl_device_binaries MDeviceBinaries = nullptr;
   std::string MPrefix;
   include_pairs_t MIncludePairs;
 };

sycl/source/detail/kernel_compiler/kernel_compiler_sycl.cpp

@@ -323,6 +323,15 @@ std::pair<sycl_device_binaries, std::string> SYCL_JIT_to_SPIRV(
 #endif
 }
 
+void SYCL_JIT_destroy([[maybe_unused]] sycl_device_binaries Binaries) {
+#if SYCL_EXT_JIT_ENABLE
+  sycl::detail::jit_compiler::get_instance().destroyDeviceBinaries(Binaries);
+#else
+  throw sycl::exception(sycl::errc::invalid,
+                        "kernel_compiler via sycl-jit is not available");
+#endif
+}
+
 } // namespace detail
 } // namespace ext::oneapi::experimental
 } // namespace _V1

sycl/source/detail/kernel_compiler/kernel_compiler_sycl.hpp

@@ -40,6 +40,8 @@ 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);
 
+void SYCL_JIT_destroy(sycl_device_binaries Binaries);
+
 bool SYCL_JIT_Compilation_Available();
 
 } // namespace detail

sycl/source/detail/program_manager/program_manager.cpp

@@ -1998,6 +1998,89 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
   }
 }
 
+void ProgramManager::removeImages(sycl_device_binaries DeviceBinary) {
+  for (int I = 0; I < DeviceBinary->NumDeviceBinaries; I++) {
+    sycl_device_binary RawImg = &(DeviceBinary->DeviceBinaries[I]);
+    const sycl_offload_entry EntriesB = RawImg->EntriesBegin;
+    const sycl_offload_entry EntriesE = RawImg->EntriesEnd;
+    // Treat the image as empty one
+    if (EntriesB == EntriesE)
+      continue;
+
+    // Retrieve RTDeviceBinaryImage by looking up the first offload entry
+    kernel_id FirstKernelID = getSYCLKernelID(RawImg->EntriesBegin->name);
+    auto RTDBImages = getRawDeviceImages({FirstKernelID});
+    assert(RTDBImages.size() == 1);
+
+    RTDeviceBinaryImage *Img = *RTDBImages.begin();
+
+    // Drop the kernel argument mask map
+    m_EliminatedKernelArgMasks.erase(Img);
+
+    // Acquire lock to modify maps for kernel bundles
+    std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
+
+    // Unmap the unique kernel IDs for the offload entries
+    for (sycl_offload_entry EntriesIt = EntriesB; EntriesIt != EntriesE;
+         ++EntriesIt) {
+
+      // Drop entry for service kernel
+      if (std::strstr(EntriesIt->name, "__sycl_service_kernel__")) {
+        m_ServiceKernels.erase(EntriesIt->name);
+        continue;
+      }
+
+      // Exported device functions won't have a kernel ID
+      if (m_ExportedSymbolImages.find(EntriesIt->name) !=
+          m_ExportedSymbolImages.end()) {
+        continue;
+      }
+
+      auto It = m_KernelName2KernelIDs.find(EntriesIt->name);
+      assert(It != m_KernelName2KernelIDs.end());
+      m_KernelName2KernelIDs.erase(It);
+      m_KernelIDs2BinImage.erase(It->second);
+    }
+
+    // Drop reverse mapping
+    m_BinImg2KernelIDs.erase(Img);
+
+    // Unregister exported symbols (needs to happen after the ID unmap loop)
+    for (const sycl_device_binary_property &ESProp :
+         Img->getExportedSymbols()) {
+      m_ExportedSymbolImages.erase(ESProp->Name);
+    }
+
+    // TODO: Handle other runtime info that was set up by `addImages`
+    assert(Img->getVirtualFunctions().empty());
+    assert(Img->getAssertUsed().empty());
+    assert(!Img->getProperty("sanUsed"));
+    assert(Img->getImplicitLocalArg().empty());
+    assert(Img->getDeviceGlobals().empty());
+    assert(Img->getHostPipes().empty());
+
+    // Purge references to the image in native programs map
+    {
+      std::lock_guard<std::mutex> NativeProgramsGuard(MNativeProgramsMutex);
+
+      // The map does not keep references to program handles; we can erase the
+      // entry without calling UR release
+      for (auto It = NativePrograms.begin(); It != NativePrograms.end();) {
+        auto CurIt = It++;
+        if (CurIt->second == Img)
+          NativePrograms.erase(CurIt);
+      }
+    }
+
+    // Finally, destroy the image by erasing the associated unique ptr
+    auto It =
+        std::find_if(m_DeviceImages.begin(), m_DeviceImages.end(),
+                     [Img](const auto &UPtr) { return UPtr.get() == Img; });
+    assert(It != m_DeviceImages.end());
+    m_DeviceImages.erase(It);
+  }
+}
+
 void ProgramManager::debugPrintBinaryImages() const {
   for (const auto &ImgIt : m_BinImg2KernelIDs) {
     ImgIt.first->print();

sycl/source/detail/program_manager/program_manager.hpp

@@ -216,6 +216,7 @@ class ProgramManager {
                                                const ContextImplPtr Context);
 
   void addImages(sycl_device_binaries DeviceImages);
+  void removeImages(sycl_device_binaries DeviceImages);
   void debugPrintBinaryImages() const;
   static std::string getProgramBuildLog(const ur_program_handle_t &Program,
                                         const ContextImplPtr &Context);