Fix issue of adoptive binaries and deallocation

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/source/detail/device_image_impl.hpp

@@ -115,6 +115,22 @@ struct KernelCompilerBinaryInfo {
   include_pairs_t MIncludePairs;
 };
 
+// Helper class to unregister shared SYCL binaries.
+class ManagedDeviceBinaries {
+public:
+  ManagedDeviceBinaries(sycl_device_binaries &&Binaries)
+      : MBinaries{Binaries} {}
+  ~ManagedDeviceBinaries() try {
+    ProgramManager::getInstance().removeImages(MBinaries);
+    syclex::detail::SYCL_JIT_Destroy(MBinaries);
+  } catch (std::exception &e) {
+    __SYCL_REPORT_EXCEPTION_TO_STREAM(
+        "exception during unregistration of SYCL binaries", e);
+  }
+private:
+  sycl_device_binaries MBinaries;
+};
+
 // The class is impl counterpart for sycl::device_image
 // It can represent a program in different states, kernel_id's it has and state
 // of specialization constants for it
@@ -550,12 +566,12 @@ class device_image_impl {
     return MRTCBinInfo && MRTCBinInfo->MLanguage == Lang;
   }
 
-  std::vector<std::shared_ptr<device_image_impl>>
-  buildFromSource(const std::vector<device> Devices,
-                  const std::vector<std::string> &BuildOptions,
-                  std::string *LogPtr,
-                  const std::vector<std::string> &RegisteredKernelNames,
-                  std::vector<sycl_device_binaries> &OutDeviceBinaries) const {
+  std::vector<std::shared_ptr<device_image_impl>> buildFromSource(
+      const std::vector<device> Devices,
+      const std::vector<std::string> &BuildOptions, std::string *LogPtr,
+      const std::vector<std::string> &RegisteredKernelNames,
+      std::vector<std::shared_ptr<ManagedDeviceBinaries>> &OutDeviceBinaries)
+      const {
     assert(!std::holds_alternative<const RTDeviceBinaryImage *>(MBinImage));
     assert(MRTCBinInfo);
     assert(MOrigins & ImageOriginKernelCompiler);
@@ -666,7 +682,8 @@ class device_image_impl {
         Result.push_back(getSyclObjImpl(ImgWithDeps.getMain()));
       }
 
-      OutDeviceBinaries.push_back(std::move(Binaries));
+      OutDeviceBinaries.emplace_back(
+          std::make_shared<ManagedDeviceBinaries>(std::move(Binaries)));
       return Result;
     }
 

sycl/source/detail/kernel_bundle_impl.hpp

@@ -302,9 +302,11 @@ class kernel_bundle_impl {
     }
 
     for (const detail::KernelBundleImplPtr &Bundle : Bundles) {
-
       MDeviceImages.insert(MDeviceImages.end(), Bundle->MDeviceImages.begin(),
                            Bundle->MDeviceImages.end());
+      MDeviceBinaries.insert(MDeviceBinaries.end(),
+                             Bundle->MDeviceBinaries.begin(),
+                             Bundle->MDeviceBinaries.end());
     }
 
     fillUniqueDeviceImages();
@@ -368,9 +370,10 @@ class kernel_bundle_impl {
 
   // oneapi_ext_kernel_compiler
   // construct from built source files
-  kernel_bundle_impl(const context &Context, const std::vector<device> &Devs,
-                     std::vector<device_image_plain> &&DevImgs,
-                     std::vector<sycl_device_binaries> &&DevBinaries)
+  kernel_bundle_impl(
+      const context &Context, const std::vector<device> &Devs,
+      std::vector<device_image_plain> &&DevImgs,
+      std::vector<std::shared_ptr<ManagedDeviceBinaries>> &&DevBinaries)
       : MContext(Context), MDevices(Devs),
         MUniqueDeviceImages(std::move(DevImgs)),
         MState(bundle_state::executable),
@@ -393,7 +396,7 @@ class kernel_bundle_impl {
     assert(allSourceBasedImages() && "All images must be source-based.");
 
     std::vector<device_image_plain> NewDevImgs;
-    std::vector<sycl_device_binaries> NewDeviceBinaries;
+    std::vector<std::shared_ptr<ManagedDeviceBinaries>> NewDeviceBinaries;
     for (device_image_plain &DevImg : MUniqueDeviceImages) {
       std::vector<std::shared_ptr<device_image_impl>> NewDevImgImpls =
           getSyclObjImpl(DevImg)->buildFromSource(Devices, BuildOptions, LogPtr,
@@ -626,17 +629,6 @@ class kernel_bundle_impl {
     return true;
   }
 
-  ~kernel_bundle_impl() {
-    try {
-      for (sycl_device_binaries &Binaries : MDeviceBinaries) {
-        ProgramManager::getInstance().removeImages(Binaries);
-        syclex::detail::SYCL_JIT_Destroy(Binaries);
-      }
-    } catch (std::exception &e) {
-      __SYCL_REPORT_EXCEPTION_TO_STREAM("exception in ~kernel_bundle_impl", e);
-    }
-  }
-
   bool hasSourceBasedImages() const noexcept {
     return std::any_of(begin(), end(), [](const device_image_plain &DevImg) {
       return getSyclObjImpl(DevImg)->getOriginMask() &
@@ -771,8 +763,8 @@ class kernel_bundle_impl {
   bundle_state MState;
 
   // For sycl_jit, building from source may have produced sycl binaries that
-  // this kernel_bundle now manages.
-  std::vector<sycl_device_binaries> MDeviceBinaries;
+  // the kernel_bundles now manage.
+  std::vector<std::shared_ptr<ManagedDeviceBinaries>> MDeviceBinaries;
 };
 
 } // namespace detail

sycl/test-e2e/KernelCompiler/sycl_join.cpp

@@ -121,12 +121,12 @@ int main() {
     sycl::free(IntPtr, Q);
   }
 
-  exe_kb KBExe1 = syclex::build(KBSrc1);
-  exe_kb KBExe2 = syclex::build(KBSrc2);
+  auto KBExe1 = std::make_shared<exe_kb>(syclex::build(KBSrc1));
+  auto KBExe2 = std::make_shared<exe_kb>(syclex::build(KBSrc2));
 
   // Test joining of source-based executable kernel bundles.
   {
-    std::vector<exe_kb> KBExes{KBExe1, KBExe2};
+    std::vector<exe_kb> KBExes{*KBExe1, *KBExe2};
 
     exe_kb KBExeJoined = sycl::join(KBExes);
     assert(KBExeJoined.ext_oneapi_has_kernel("TestKernel1"));
@@ -173,7 +173,7 @@ int main() {
     exe_kb RegularKBExe =
         sycl::get_kernel_bundle<sycl::bundle_state::executable>(
             Ctx, RegularSYCLKernelIDs);
-    std::vector<exe_kb> KBExes{KBExe1, KBExe2, RegularKBExe};
+    std::vector<exe_kb> KBExes{*KBExe1, *KBExe2, RegularKBExe};
 
     exe_kb KBExeJoined = sycl::join(KBExes);
     assert(KBExeJoined.ext_oneapi_has_kernel("TestKernel1"));
@@ -235,5 +235,50 @@ int main() {
     sycl::free(IntPtr, Q);
   }
 
+  // Test joining of executable kernel bundles with the original bundles dying
+  // before the parent.
+  {
+    std::vector<exe_kb> KBExes{*KBExe1, *KBExe2};
+
+    KBExe1.reset();
+    KBExe2.reset();
+
+    exe_kb KBExeJoined = sycl::join(KBExes);
+    assert(KBExeJoined.ext_oneapi_has_kernel("TestKernel1"));
+    assert(KBExeJoined.ext_oneapi_has_kernel("TestKernel2"));
+
+    sycl::kernel K1 = KBExeJoined.ext_oneapi_get_kernel("TestKernel1");
+    sycl::kernel K2 = KBExeJoined.ext_oneapi_get_kernel("TestKernel2");
+
+    int *IntPtr = sycl::malloc_shared<int>(1, Q);
+    *IntPtr = 0;
+
+    Q.submit([&](sycl::handler &CGH) {
+       CGH.set_args(IntPtr);
+       CGH.single_task(K1);
+     }).wait_and_throw();
+
+    if (*IntPtr != 42) {
+      std::cout << "TestKernel1 in joined source-based executable bundles with "
+                   "dead parents failed: "
+                << *IntPtr << " != 42\n";
+      ++Failed;
+    }
+
+    Q.submit([&](sycl::handler &CGH) {
+       CGH.set_args(IntPtr);
+       CGH.single_task(K2);
+     }).wait_and_throw();
+
+    if (*IntPtr != 24) {
+      std::cout << "TestKernel1 in joined source-based executable bundles with "
+                   "dead parents failed: "
+                << *IntPtr << " != 24\n";
+      ++Failed;
+    }
+
+    sycl::free(IntPtr, Q);
+  }
+
   return Failed;
 }