Untangle destruction of device global map entries

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

Author: jopperm

sycl/source/detail/context_impl.cpp

@@ -329,6 +329,11 @@ void context_impl::addAssociatedDeviceGlobal(const void *DeviceGlobalPtr) {
   MAssociatedDeviceGlobals.insert(DeviceGlobalPtr);
 }
 
+void context_impl::removeAssociatedDeviceGlobal(const void *DeviceGlobalPtr) {
+  std::lock_guard<std::mutex> Lock{MAssociatedDeviceGlobalsMutex};
+  MAssociatedDeviceGlobals.erase(DeviceGlobalPtr);
+}
+
 void context_impl::addDeviceGlobalInitializer(
     ur_program_handle_t Program, const std::vector<device> &Devs,
     const RTDeviceBinaryImage *BinImage) {
@@ -407,6 +412,9 @@ std::vector<ur_event_handle_t> context_impl::initializeDeviceGlobals(
     // Device global map entry pointers will not die before the end of the
     // program and the pointers will stay the same, so we do not need
     // m_DeviceGlobalsMutex here.
+    // The lifetimes of device global map entries representing globals in
+    // runtime-compiled code will be tied to the kernel bundle, so the
+    // assumption holds in that setting as well.
     for (DeviceGlobalMapEntry *DeviceGlobalEntry : DeviceGlobalEntries) {
       // Get or allocate the USM memory associated with the device global.
       DeviceGlobalUSMMem &DeviceGlobalUSM =

sycl/source/detail/context_impl.hpp

@@ -203,6 +203,9 @@ class context_impl {
   /// Adds an associated device global to the tracked associates.
   void addAssociatedDeviceGlobal(const void *DeviceGlobalPtr);
 
+  /// Removes an associated device global from the tracked associates.
+  void removeAssociatedDeviceGlobal(const void *DeviceGlobalPtr);
+
   /// Adds a device global initializer.
   void addDeviceGlobalInitializer(ur_program_handle_t Program,
                                   const std::vector<device> &Devs,

sycl/source/detail/kernel_bundle_impl.hpp

@@ -385,6 +385,7 @@ class kernel_bundle_impl {
       std::vector<std::string> &&KernelNames,
       std::unordered_map<std::string, std::string> &&MangledKernelNames,
       std::vector<std::string> &&DeviceGlobalNames,
+      std::vector<std::unique_ptr<std::byte[]>> &&DeviceGlobalAllocations,
       sycl_device_binaries Binaries, std::string &&Prefix,
       syclex::source_language Lang)
       : kernel_bundle_impl(std::move(Ctx), std::move(Devs), KernelIDs,
@@ -399,6 +400,7 @@ class kernel_bundle_impl {
     MKernelNames = std::move(KernelNames);
     MMangledKernelNames = std::move(MangledKernelNames);
     MDeviceGlobalNames = std::move(DeviceGlobalNames);
+    MDeviceGlobalAllocations = std::move(DeviceGlobalAllocations);
     MDeviceBinaries = Binaries;
     MPrefix = std::move(Prefix);
     MLanguage = Lang;
@@ -535,6 +537,12 @@ class kernel_bundle_impl {
       std::vector<kernel_id> KernelIDs;
       std::vector<std::string> KernelNames;
       std::unordered_map<std::string, std::string> MangledKernelNames;
+
+      std::unordered_set<std::string> DeviceGlobalIDSet;
+      std::vector<std::string> DeviceGlobalIDVec;
+      std::vector<std::string> DeviceGlobalNames;
+      std::vector<std::unique_ptr<std::byte[]>> DeviceGlobalAllocations;
+
       for (const auto &KernelID : PM.getAllSYCLKernelIDs()) {
         std::string_view KernelName{KernelID.get_name()};
         if (KernelName.find(Prefix) == 0) {
@@ -552,8 +560,8 @@ class kernel_bundle_impl {
         }
       }
 
-      // Apply frontend information.
       for (const auto *RawImg : PM.getRawDeviceImages(KernelIDs)) {
+        // Mangled names.
         for (const sycl_device_binary_property &RKProp :
              RawImg->getRegisteredKernels()) {
 
@@ -563,14 +571,8 @@ class kernel_bundle_impl {
               reinterpret_cast<const char *>(BA.begin()), MangledNameLen};
           MangledKernelNames.emplace(RKProp->Name, MangledName);
         }
-      }
 
-      // Determine IDs of all device globals referenced by this bundle's
-      // kernels. These IDs are also prefixed.
-      std::unordered_set<std::string> DeviceGlobalIDSet;
-      std::vector<std::string> DeviceGlobalIDVec;
-      std::vector<std::string> DeviceGlobalNames;
-      for (const auto &RawImg : PM.getRawDeviceImages(KernelIDs)) {
+        // Device globals.
         for (const auto &DeviceGlobalProp : RawImg->getDeviceGlobals()) {
           std::string_view DeviceGlobalName{DeviceGlobalProp->Name};
           assert(DeviceGlobalName.find(Prefix) == 0);
@@ -585,12 +587,6 @@ class kernel_bundle_impl {
         }
       }
 
-      // Create the executable bundle.
-      auto ExecBundle = std::make_shared<kernel_bundle_impl>(
-          MContext, MDevices, KernelIDs, std::move(KernelNames),
-          std::move(MangledKernelNames), std::move(DeviceGlobalNames), Binaries,
-          std::move(Prefix), MLanguage);
-
       // Device globals are usually statically allocated and registered in the
       // integration footer, which we don't have in the RTC context. Instead, we
       // dynamically allocate storage tied to the executable kernel bundle.
@@ -599,13 +595,13 @@ class kernel_bundle_impl {
 
         size_t AllocSize = DeviceGlobalEntry->MDeviceGlobalTSize; // init value
         if (!DeviceGlobalEntry->MIsDeviceImageScopeDecorated) {
-          // USM pointer. TODO: it's actually a decorated multi_ptr.
+          // Consider storage for device USM pointer.
           AllocSize += sizeof(void *);
         }
         auto Alloc = std::make_unique<std::byte[]>(AllocSize);
         std::string_view DeviceGlobalName{DeviceGlobalEntry->MUniqueId};
         PM.addOrInitDeviceGlobalEntry(Alloc.get(), DeviceGlobalName.data());
-        ExecBundle->MDeviceGlobalAllocations.push_back(std::move(Alloc));
+        DeviceGlobalAllocations.push_back(std::move(Alloc));
 
         // Drop the RTC prefix from the entry's symbol name. Note that the PM
         // still manages this device global under its prefixed name.
@@ -614,7 +610,11 @@ class kernel_bundle_impl {
         DeviceGlobalEntry->MUniqueId = DeviceGlobalName;
       }
 
-      return ExecBundle;
+      return std::make_shared<kernel_bundle_impl>(
+          MContext, MDevices, KernelIDs, std::move(KernelNames),
+          std::move(MangledKernelNames), std::move(DeviceGlobalNames),
+          std::move(DeviceGlobalAllocations), Binaries, std::move(Prefix),
+          MLanguage);
     }
 
     ur_program_handle_t UrProgram = nullptr;
@@ -781,6 +781,28 @@ class kernel_bundle_impl {
     return Entries.front();
   }
 
+  void unregister_device_globals_from_context() {
+    if (MDeviceGlobalNames.empty())
+      return;
+
+    // Manually trigger the release of resources for all device global map
+    // entries associated with this runtime-compiled bundle. Normally, this
+    // would happen in `~context_impl()`, however in the RTC setting, the
+    // context outlives the DG map entries owned by the program manager.
+
+    std::vector<std::string> DeviceGlobalIDs;
+    std::transform(MDeviceGlobalNames.begin(), MDeviceGlobalNames.end(),
+                   std::back_inserter(DeviceGlobalIDs),
+                   [&](const std::string &DGName) { return MPrefix + DGName; });
+    auto ContextImpl = getSyclObjImpl(MContext);
+    for (DeviceGlobalMapEntry *Entry :
+         ProgramManager::getInstance().getDeviceGlobalEntries(
+             DeviceGlobalIDs)) {
+      Entry->removeAssociatedResources(ContextImpl.get());
+      ContextImpl->removeAssociatedDeviceGlobal(Entry->MDeviceGlobalPtr);
+    }
+  }
+
 public:
   bool ext_oneapi_has_kernel(const std::string &Name) {
     return is_kernel_name(adjust_kernel_name(Name));
@@ -1121,6 +1143,7 @@ class kernel_bundle_impl {
   ~kernel_bundle_impl() {
     try {
       if (MDeviceBinaries) {
+        unregister_device_globals_from_context();
         ProgramManager::getInstance().removeImages(MDeviceBinaries);
         syclex::detail::SYCL_JIT_destroy(MDeviceBinaries);
       }
@@ -1162,11 +1185,10 @@ class kernel_bundle_impl {
   std::vector<std::string> MKernelNames;
   std::unordered_map<std::string, std::string> MMangledKernelNames;
   std::vector<std::string> MDeviceGlobalNames;
+  std::vector<std::unique_ptr<std::byte[]>> MDeviceGlobalAllocations;
   sycl_device_binaries MDeviceBinaries = nullptr;
   std::string MPrefix;
   include_pairs_t MIncludePairs;
-
-  std::vector<std::unique_ptr<std::byte[]>> MDeviceGlobalAllocations;
 };
 
 } // namespace detail