[Clang][SYCL] Discard pointers to free function descriptors when library is unloaded

Author: uditagarwal97

clang/lib/Sema/SemaSYCL.cpp

@@ -7313,6 +7313,11 @@ void SYCLIntegrationHeader::emit(raw_ostream &O) {
       << "sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, "
       << KernelDescs.size() << ");\n";
     O << "  }\n";
+    O << "  ~GlobalMapUpdater() {\n";
+    O << "    sycl::detail::free_function_info_map::remove("
+      << "sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, "
+      << KernelDescs.size() << ");\n";
+    O << "  }\n";
     O << "};\n";
     O << "static GlobalMapUpdater updater;\n";
     O << "} // namespace detail\n";

clang/test/CodeGenSYCL/free_function_default_template_arguments.cpp

@@ -1094,6 +1094,9 @@ namespace Testing::Tests {
 // CHECK-NEXT:   GlobalMapUpdater() {
 // CHECK-NEXT:     sycl::detail::free_function_info_map::add(sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, 28);
 // CHECK-NEXT:   }
+// CHECK-NEXT:   ~GlobalMapUpdater() {
+// CHECK-NEXT:     sycl::detail::free_function_info_map::remove(sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, 28);
+// CHECK-NEXT:   }
 // CHECK-NEXT: };
 // CHECK-NEXT: static GlobalMapUpdater updater;
 // CHECK-NEXT: } // namespace detail

clang/test/CodeGenSYCL/free_function_int_header.cpp

@@ -1574,6 +1574,9 @@ void ff_24(int arg) {
 // CHECK-NEXT:  GlobalMapUpdater() {
 // CHECK-NEXT:     sycl::detail::free_function_info_map::add(sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, 33);
 // CHECK-NEXT:   }
+// CHECK-NEXT:  ~GlobalMapUpdater() {
+// CHECK-NEXT:     sycl::detail::free_function_info_map::remove(sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, 33);
+// CHECK-NEXT:   }
 // CHECK-NEXT: };
 // CHECK-NEXT: static GlobalMapUpdater updater;
 // CHECK-NEXT: } // namespace detail

clang/test/CodeGenSYCL/free_function_int_header_rtc_mode.cpp

@@ -82,6 +82,9 @@ int main(){
 // CHECK-NORTC-NEXT:   GlobalMapUpdater() {
 // CHECK-NORTC-NEXT:     sycl::detail::free_function_info_map::add(sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, 3);
 // CHECK-NORTC-NEXT:   }
+// CHECK-NORTC-NEXT:   ~GlobalMapUpdater() {
+// CHECK-NORTC-NEXT:     sycl::detail::free_function_info_map::remove(sycl::detail::kernel_names, sycl::detail::kernel_args_sizes, 3);
+// CHECK-NORTC-NEXT:   }
 // CHECK-NORTC-NEXT: };
 // CHECK-NORTC-NEXT: static GlobalMapUpdater updater;
 // CHECK-NORTC-NEXT: } // namespace detail

clang/test/SemaSYCL/Inputs/sycl/detail/kernel_global_info.hpp

@@ -16,6 +16,7 @@ namespace detail {
 namespace free_function_info_map {
 
 __SYCL_EXPORT void add(const void *DeviceGlobalPtr, const char *UniqueId);
+__SYCL_EXPORT void remove(const void *DeviceGlobalPtr, const char *UniqueId);
 
 } // namespace free_function_info_map
 } // namespace detail

sycl/include/sycl/detail/kernel_global_info.hpp

@@ -18,6 +18,8 @@ namespace free_function_info_map {
 __SYCL_EXPORT void add(const char *const *UniqueId,
                        const unsigned *DeviceGlobalPtr, unsigned Size);
 
+__SYCL_EXPORT void remove(const char *const *UniqueId,
+                          const unsigned *DeviceGlobalPtr, unsigned Size);
 } // namespace free_function_info_map
 } // namespace detail
 } // namespace _V1

sycl/source/detail/kernel_global_info.cpp

@@ -23,6 +23,17 @@ __SYCL_EXPORT void add(const char *const *FreeFunctionNames,
       std::move(GlobalInfoToCopy));
 }
 
+__SYCL_EXPORT void remove(const char *const *FreeFunctionNames,
+                          const unsigned *FreeFunctionNumArgs, unsigned Size) {
+  std::unordered_map<std::string_view, unsigned> GlobalInfoToCopy;
+  for (size_t i = 0; i < Size; ++i) {
+    GlobalInfoToCopy[std::string_view{FreeFunctionNames[i]}] =
+        FreeFunctionNumArgs[i];
+  }
+  detail::ProgramManager::getInstance().unRegisterKernelGlobalInfo(
+      std::move(GlobalInfoToCopy));
+}
+
 } // namespace detail::free_function_info_map
 } // namespace _V1
 } // namespace sycl

sycl/source/detail/program_manager/program_manager.cpp

@@ -2417,12 +2417,22 @@ void ProgramManager::registerKernelGlobalInfo(
   if (m_FreeFunctionKernelGlobalInfo.empty())
     m_FreeFunctionKernelGlobalInfo = std::move(GlobalInfoToCopy);
   else {
-    for (auto &GlobalInfo : GlobalInfoToCopy) {
+    for (auto &GlobalInfo : GlobalInfoToCopy)
       m_FreeFunctionKernelGlobalInfo.insert(GlobalInfo);
-    }
   }
 }
 
+// Remove entries from m_FreeFunctionKernelGlobalInfo that matches
+// the ones in GlobalInfoToCopy. This function is called when a shared
+// library consisting of SYCL kernels is unloaded.
+void ProgramManager::unRegisterKernelGlobalInfo(
+    std::unordered_map<std::string_view, unsigned> &&GlobalInfoToCopy) {
+  std::lock_guard<std::mutex> Guard(MNativeProgramsMutex);
+
+  for (const auto &GlobalInfo : GlobalInfoToCopy)
+    m_FreeFunctionKernelGlobalInfo.erase(GlobalInfo.first);
+}
+
 std::optional<unsigned>
 ProgramManager::getKernelGlobalInfoDesc(const char *UniqueId) {
   std::lock_guard<std::mutex> Guard(MNativeProgramsMutex);

sycl/source/detail/program_manager/program_manager.hpp

@@ -259,6 +259,12 @@ class ProgramManager {
   void registerKernelGlobalInfo(
       std::unordered_map<std::string_view, unsigned> &&GlobalInfoToCopy);
 
+  // The function removes kernel global descriptors from the
+  // kernel global map when a shared library consisting o SYCL kernels
+  // is unloaded.
+  void unRegisterKernelGlobalInfo(
+      std::unordered_map<std::string_view, unsigned> &&GlobalInfoToCopy);
+
   // The function returns a pointer to the kernel global desc identified by
   // the unique ID from the kernel global map.
   std::optional<unsigned> getKernelGlobalInfoDesc(const char *UniqueId);

sycl/test/abi/sycl_symbols_linux.dump

@@ -3344,6 +3344,7 @@ _ZN4sycl3_V16detail21LocalAccessorBaseHostC1ENS0_5rangeILi3EEEiiRKNS0_13property
 _ZN4sycl3_V16detail21LocalAccessorBaseHostC2ENS0_5rangeILi3EEEiiRKNS0_13property_listE
 _ZN4sycl3_V16detail22addHostAccessorAndWaitEPNS1_16AccessorImplHostE
 _ZN4sycl3_V16detail22free_function_info_map3addEPKPKcPKjj
+_ZN4sycl3_V16detail22free_function_info_map6removeEPKPKcPKjj
 _ZN4sycl3_V16detail22get_kernel_bundle_implERKNS0_7contextERKSt6vectorINS0_6deviceESaIS6_EENS0_12bundle_stateE
 _ZN4sycl3_V16detail22get_kernel_bundle_implERKNS0_7contextERKSt6vectorINS0_6deviceESaIS6_EENS0_12bundle_stateERKSt8functionIFbRKSt10shared_ptrINS1_17device_image_implEEEE
 _ZN4sycl3_V16detail22get_kernel_bundle_implERKNS0_7contextERKSt6vectorINS0_6deviceESaIS6_EERKNS0_4spanIcLm18446744073709551615EEENS0_12bundle_stateE