[NFC][SYCL] std::enable_shared_from_this for device_image_impl (#19182)

aelovikov-intel · web-flow · commit 6eb4a923c83a · 2025-06-26T16:50:00.000-07:00
diff --git a/sycl/source/backend.cpp b/sycl/source/backend.cpp
@@ -301,9 +301,9 @@ make_kernel_bundle(ur_native_handle_t NativeHandle,
   // do the same to user images, since they may contain references to undefined
   // symbols (e.g. when kernel_bundle is supposed to be joined with another).
   auto KernelIDs = std::make_shared<std::vector<kernel_id>>();
-  auto DevImgImpl = std::make_shared<device_image_impl>(
-      nullptr, TargetContext, Devices, State, KernelIDs, UrProgram,
-      ImageOriginInterop);
+  auto DevImgImpl =
+      device_image_impl::create(nullptr, TargetContext, Devices, State,
+                                KernelIDs, UrProgram, ImageOriginInterop);
   device_image_plain DevImg{DevImgImpl};
 
   return kernel_bundle_impl::create(TargetContext, Devices, DevImg);
diff --git a/sycl/source/detail/device_image_impl.cpp b/sycl/source/detail/device_image_impl.cpp
@@ -15,8 +15,7 @@ namespace detail {
 
 std::shared_ptr<kernel_impl> device_image_impl::tryGetExtensionKernel(
     std::string_view Name, const context &Context,
-    const kernel_bundle_impl &OwnerBundle,
-    const std::shared_ptr<device_image_impl> &Self) const {
+    const kernel_bundle_impl &OwnerBundle) {
   if (!(getOriginMask() & ImageOriginKernelCompiler) &&
       !((getOriginMask() & ImageOriginSYCLBIN) && hasKernelName(Name)))
     return nullptr;
@@ -34,9 +33,9 @@ std::shared_ptr<kernel_impl> device_image_impl::tryGetExtensionKernel(
       auto [UrKernel, CacheMutex, ArgMask] =
           PM.getOrCreateKernel(Context, AdjustedName,
                                /*PropList=*/{}, UrProgram);
-      return std::make_shared<kernel_impl>(UrKernel, *getSyclObjImpl(Context),
-                                           Self, OwnerBundle.shared_from_this(),
-                                           ArgMask, UrProgram, CacheMutex);
+      return std::make_shared<kernel_impl>(
+          UrKernel, *getSyclObjImpl(Context), shared_from_this(),
+          OwnerBundle.shared_from_this(), ArgMask, UrProgram, CacheMutex);
     }
     return nullptr;
   }
@@ -49,7 +48,7 @@ std::shared_ptr<kernel_impl> device_image_impl::tryGetExtensionKernel(
   // Kernel created by urKernelCreate is implicitly retained.
 
   return std::make_shared<kernel_impl>(
-      UrKernel, *detail::getSyclObjImpl(Context), Self,
+      UrKernel, *detail::getSyclObjImpl(Context), shared_from_this(),
       OwnerBundle.shared_from_this(),
       /*ArgMask=*/nullptr, UrProgram, /*CacheMutex=*/nullptr);
 }
diff --git a/sycl/source/detail/device_image_impl.hpp b/sycl/source/detail/device_image_impl.hpp
@@ -230,7 +230,12 @@ struct KernelCompilerBinaryInfo {
 // 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
-class device_image_impl {
+class device_image_impl
+    : public std::enable_shared_from_this<device_image_impl> {
+  struct private_tag {
+    explicit private_tag() = default;
+  };
+
 public:
   // The struct maps specialization ID to offset in the binary blob where value
   // for this spec const should be.
@@ -249,8 +254,7 @@ class device_image_impl {
   device_image_impl(const RTDeviceBinaryImage *BinImage, context Context,
                     std::vector<device> Devices, bundle_state State,
                     std::shared_ptr<std::vector<kernel_id>> KernelIDs,
-                    ur_program_handle_t Program,
-                    uint8_t Origins = ImageOriginSYCLOffline)
+                    ur_program_handle_t Program, uint8_t Origins, private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
         MKernelIDs(std::move(KernelIDs)),
@@ -272,7 +276,7 @@ class device_image_impl {
       const std::vector<unsigned char> &SpecConstsBlob, uint8_t Origins,
       std::optional<KernelCompilerBinaryInfo> &&RTCInfo,
       KernelNameSetT &&KernelNames,
-      std::unique_ptr<DynRTDeviceBinaryImage> &&MergedImageStorage = nullptr)
+      std::unique_ptr<DynRTDeviceBinaryImage> &&MergedImageStorage, private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
         MKernelIDs(std::move(KernelIDs)), MKernelNames{std::move(KernelNames)},
@@ -285,7 +289,7 @@ class device_image_impl {
   device_image_impl(const RTDeviceBinaryImage *BinImage, const context &Context,
                     const std::vector<device> &Devices, bundle_state State,
                     ur_program_handle_t Program, syclex::source_language Lang,
-                    KernelNameSetT &&KernelNames)
+                    KernelNameSetT &&KernelNames, private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
         MKernelIDs(std::make_shared<std::vector<kernel_id>>()),
@@ -302,7 +306,8 @@ class device_image_impl {
       std::shared_ptr<std::vector<kernel_id>> &&KernelIDs,
       syclex::source_language Lang, KernelNameSetT &&KernelNames,
       MangledKernelNameMapT &&MangledKernelNames, std::string &&Prefix,
-      std::shared_ptr<ManagedDeviceGlobalsRegistry> &&DeviceGlobalRegistry)
+      std::shared_ptr<ManagedDeviceGlobalsRegistry> &&DeviceGlobalRegistry,
+      private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(nullptr),
         MKernelIDs(std::move(KernelIDs)), MKernelNames{std::move(KernelNames)},
@@ -317,7 +322,7 @@ class device_image_impl {
   device_image_impl(const std::string &Src, context Context,
                     const std::vector<device> &Devices,
                     syclex::source_language Lang,
-                    include_pairs_t &&IncludePairsVec)
+                    include_pairs_t &&IncludePairsVec, private_tag)
       : MBinImage(Src), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(bundle_state::ext_oneapi_source),
         MProgram(nullptr),
@@ -331,7 +336,7 @@ class device_image_impl {
 
   device_image_impl(const std::vector<std::byte> &Bytes, const context &Context,
                     const std::vector<device> &Devices,
-                    syclex::source_language Lang)
+                    syclex::source_language Lang, private_tag)
       : MBinImage(Bytes), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(bundle_state::ext_oneapi_source),
         MProgram(nullptr),
@@ -344,7 +349,8 @@ class device_image_impl {
 
   device_image_impl(const context &Context, const std::vector<device> &Devices,
                     bundle_state State, ur_program_handle_t Program,
-                    syclex::source_language Lang, KernelNameSetT &&KernelNames)
+                    syclex::source_language Lang, KernelNameSetT &&KernelNames,
+                    private_tag)
       : MBinImage(static_cast<const RTDeviceBinaryImage *>(nullptr)),
         MContext(std::move(Context)), MDevices(std::move(Devices)),
         MState(State), MProgram(Program),
@@ -354,6 +360,12 @@ class device_image_impl {
         MOrigins(ImageOriginKernelCompiler),
         MRTCBinInfo(KernelCompilerBinaryInfo{Lang}) {}
 
+  template <typename... Ts>
+  static std::shared_ptr<device_image_impl> create(Ts &&...args) {
+    return std::make_shared<device_image_impl>(std::forward<Ts>(args)...,
+                                               private_tag{});
+  }
+
   bool has_kernel(const kernel_id &KernelIDCand) const noexcept {
     return std::binary_search(MKernelIDs->begin(), MKernelIDs->end(),
                               KernelIDCand, LessByHash<kernel_id>{});
@@ -631,8 +643,7 @@ class device_image_impl {
 
   std::shared_ptr<kernel_impl>
   tryGetExtensionKernel(std::string_view Name, const context &Context,
-                        const kernel_bundle_impl &OwnerBundle,
-                        const std::shared_ptr<device_image_impl> &Self) const;
+                        const kernel_bundle_impl &OwnerBundle);
 
   bool hasDeviceGlobalName(const std::string &Name) const noexcept {
     if (!MRTCBinInfo.has_value())
@@ -752,9 +763,9 @@ class device_image_impl {
           *SourceStrPtr, UrProgram);
     }
     return std::vector<std::shared_ptr<device_image_impl>>{
-        std::make_shared<device_image_impl>(
-            MContext, Devices, bundle_state::executable, UrProgram,
-            MRTCBinInfo->MLanguage, std::move(KernelNameSet))};
+        device_image_impl::create(MContext, Devices, bundle_state::executable,
+                                  UrProgram, MRTCBinInfo->MLanguage,
+                                  std::move(KernelNameSet))};
   }
 
   std::vector<std::shared_ptr<device_image_impl>> compileFromSource(
@@ -1114,7 +1125,7 @@ class device_image_impl {
 
       // Mark the image as input so the program manager will bring it into
       // the right state.
-      auto DevImgImpl = std::make_shared<device_image_impl>(
+      auto DevImgImpl = device_image_impl::create(
           NewImage, MContext, std::move(SupportingDevs), bundle_state::input,
           std::move(KernelIDs), MRTCBinInfo->MLanguage, std::move(KernelNames),
           std::move(MangledKernelNames), std::string{Prefix},
diff --git a/sycl/source/detail/kernel_bundle_impl.hpp b/sycl/source/detail/kernel_bundle_impl.hpp
@@ -528,7 +528,7 @@ class kernel_bundle_impl
                      const std::string &Src, include_pairs_t IncludePairsVec,
                      private_tag)
       : MContext(Context), MDevices(Context.get_devices()),
-        MDeviceImages{device_image_plain{std::make_shared<device_image_impl>(
+        MDeviceImages{device_image_plain{device_image_impl::create(
             Src, MContext, MDevices, Lang, std::move(IncludePairsVec))}},
         MUniqueDeviceImages{MDeviceImages[0].getMain()},
         MState(bundle_state::ext_oneapi_source) {
@@ -540,8 +540,8 @@ class kernel_bundle_impl
   kernel_bundle_impl(const context &Context, syclex::source_language Lang,
                      const std::vector<std::byte> &Bytes, private_tag)
       : MContext(Context), MDevices(Context.get_devices()),
-        MDeviceImages{device_image_plain{std::make_shared<device_image_impl>(
-            Bytes, MContext, MDevices, Lang)}},
+        MDeviceImages{device_image_plain{
+            device_image_impl::create(Bytes, MContext, MDevices, Lang)}},
         MUniqueDeviceImages{MDeviceImages[0].getMain()},
         MState(bundle_state::ext_oneapi_source) {
     common_ctor_checks();
@@ -583,7 +583,7 @@ class kernel_bundle_impl
         SYCLBIN->getBestCompatibleImages(Devs);
     MDeviceImages.reserve(BestImages.size());
     for (const detail::RTDeviceBinaryImage *Image : BestImages)
-      MDeviceImages.emplace_back(std::make_shared<detail::device_image_impl>(
+      MDeviceImages.emplace_back(device_image_impl::create(
           Image, Context, Devs, ProgramManager::getBinImageState(Image),
           /*KernelIDs=*/nullptr, /*URProgram=*/nullptr, ImageOriginSYCLBIN));
     ProgramManager::getInstance().bringSYCLDeviceImagesToState(MDeviceImages,
@@ -669,8 +669,7 @@ class kernel_bundle_impl
       const std::shared_ptr<device_image_impl> &DevImgImpl =
           getSyclObjImpl(DevImg);
       if (std::shared_ptr<kernel_impl> PotentialKernelImpl =
-              DevImgImpl->tryGetExtensionKernel(Name, MContext, *this,
-                                                DevImgImpl))
+              DevImgImpl->tryGetExtensionKernel(Name, MContext, *this))
         return detail::createSyclObjFromImpl<kernel>(
             std::move(PotentialKernelImpl));
     }
@@ -1008,8 +1007,7 @@ class kernel_bundle_impl
       const std::shared_ptr<device_image_impl> &DevImgImpl =
           getSyclObjImpl(DevImg);
       if (std::shared_ptr<kernel_impl> SourceBasedKernel =
-              DevImgImpl->tryGetExtensionKernel(Name, MContext, *this,
-                                                DevImgImpl))
+              DevImgImpl->tryGetExtensionKernel(Name, MContext, *this))
         return SourceBasedKernel;
     }
 
diff --git a/sycl/source/detail/program_manager/program_manager.cpp b/sycl/source/detail/program_manager/program_manager.cpp
@@ -838,14 +838,13 @@ ProgramManager::collectDependentDeviceImagesForVirtualFunctions(
   return DeviceImagesToLink;
 }
 
-static void
-setSpecializationConstants(const std::shared_ptr<device_image_impl> &InputImpl,
-                           ur_program_handle_t Prog,
-                           const AdapterPtr &Adapter) {
-  std::lock_guard<std::mutex> Lock{InputImpl->get_spec_const_data_lock()};
+static void setSpecializationConstants(device_image_impl &InputImpl,
+                                       ur_program_handle_t Prog,
+                                       const AdapterPtr &Adapter) {
+  std::lock_guard<std::mutex> Lock{InputImpl.get_spec_const_data_lock()};
   const std::map<std::string, std::vector<device_image_impl::SpecConstDescT>>
-      &SpecConstData = InputImpl->get_spec_const_data_ref();
-  const SerializedObj &SpecConsts = InputImpl->get_spec_const_blob_ref();
+      &SpecConstData = InputImpl.get_spec_const_data_ref();
+  const SerializedObj &SpecConsts = InputImpl.get_spec_const_blob_ref();
 
   // Set all specialization IDs from descriptors in the input device image.
   for (const auto &[SpecConstNames, SpecConstDescs] : SpecConstData) {
@@ -941,7 +940,7 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
     if (!DeviceCodeWasInCache && MainImg.supportsSpecConstants()) {
       enableITTAnnotationsIfNeeded(NativePrg, Adapter);
       if (DevImgWithDeps)
-        setSpecializationConstants(getSyclObjImpl(DevImgWithDeps->getMain()),
+        setSpecializationConstants(*getSyclObjImpl(DevImgWithDeps->getMain()),
                                    NativePrg, Adapter);
     }
 
@@ -982,7 +981,7 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
           enableITTAnnotationsIfNeeded(NativePrg, Adapter);
           if (DevImgWithDeps)
             setSpecializationConstants(
-                getSyclObjImpl(DevImgWithDeps->getAll()[I]), NativePrg,
+                *getSyclObjImpl(DevImgWithDeps->getAll()[I]), NativePrg,
                 Adapter);
         }
         ProgramsToLink.push_back(NativePrg);
@@ -2508,9 +2507,9 @@ device_image_plain ProgramManager::getDeviceImageFromBinaryImage(
     KernelIDs = m_BinImg2KernelIDs[BinImage];
   }
 
-  DeviceImageImplPtr Impl = std::make_shared<detail::device_image_impl>(
+  DeviceImageImplPtr Impl = device_image_impl::create(
       BinImage, Ctx, std::vector<device>{Dev}, ImgState, KernelIDs,
-      /*PIProgram=*/nullptr);
+      /*PIProgram=*/nullptr, ImageOriginSYCLOffline);
 
   return createSyclObjFromImpl<device_image_plain>(std::move(Impl));
 }
@@ -2668,9 +2667,10 @@ ProgramManager::getSYCLDeviceImagesWithCompatibleState(
     if (ImgInfoPair.second.RequirementCounter == 0)
       continue;
 
-    DeviceImageImplPtr MainImpl = std::make_shared<detail::device_image_impl>(
+    DeviceImageImplPtr MainImpl = device_image_impl::create(
         ImgInfoPair.first, Ctx, Devs, ImgInfoPair.second.State,
-        ImgInfoPair.second.KernelIDs, /*PIProgram=*/nullptr);
+        ImgInfoPair.second.KernelIDs, /*PIProgram=*/nullptr,
+        ImageOriginSYCLOffline);
 
     std::vector<device_image_plain> Images;
     const std::set<RTDeviceBinaryImage *> &Deps = ImgInfoPair.second.Deps;
@@ -2701,8 +2701,9 @@ device_image_plain ProgramManager::createDependencyImage(
 
   assert(DepState == getBinImageState(DepImage) &&
          "State mismatch between main image and its dependency");
-  DeviceImageImplPtr DepImpl = std::make_shared<detail::device_image_impl>(
-      DepImage, Ctx, Devs, DepState, DepKernelIDs, /*PIProgram=*/nullptr);
+  DeviceImageImplPtr DepImpl =
+      device_image_impl::create(DepImage, Ctx, Devs, DepState, DepKernelIDs,
+                                /*PIProgram=*/nullptr, ImageOriginSYCLOffline);
 
   return createSyclObjFromImpl<device_image_plain>(std::move(DepImpl));
 }
@@ -2857,42 +2858,42 @@ ProgramManager::compile(const DevImgPlainWithDeps &ImgWithDeps,
   std::vector<device_image_plain> CompiledImages;
   CompiledImages.reserve(ImgWithDeps.size());
   for (const device_image_plain &DeviceImage : ImgWithDeps.getAll()) {
-    const std::shared_ptr<device_image_impl> &InputImpl =
-        getSyclObjImpl(DeviceImage);
+    device_image_impl &InputImpl = *getSyclObjImpl(DeviceImage);
 
     const AdapterPtr &Adapter =
-        getSyclObjImpl(InputImpl->get_context())->getAdapter();
+        getSyclObjImpl(InputImpl.get_context())->getAdapter();
 
     ur_program_handle_t Prog =
-        createURProgram(*InputImpl->get_bin_image_ref(),
-                        *getSyclObjImpl(InputImpl->get_context()), Devs);
+        createURProgram(*InputImpl.get_bin_image_ref(),
+                        *getSyclObjImpl(InputImpl.get_context()), Devs);
 
-    if (InputImpl->get_bin_image_ref()->supportsSpecConstants())
+    if (InputImpl.get_bin_image_ref()->supportsSpecConstants())
       setSpecializationConstants(InputImpl, Prog, Adapter);
 
-    KernelNameSetT KernelNames = InputImpl->getKernelNames();
+    KernelNameSetT KernelNames = InputImpl.getKernelNames();
 
     std::optional<detail::KernelCompilerBinaryInfo> RTCInfo =
-        InputImpl->getRTCInfo();
-    DeviceImageImplPtr ObjectImpl = std::make_shared<detail::device_image_impl>(
-        InputImpl->get_bin_image_ref(), InputImpl->get_context(),
+        InputImpl.getRTCInfo();
+    DeviceImageImplPtr ObjectImpl = device_image_impl::create(
+        InputImpl.get_bin_image_ref(), InputImpl.get_context(),
         std::vector<device>{Devs}, bundle_state::object,
-        InputImpl->get_kernel_ids_ptr(), Prog,
-        InputImpl->get_spec_const_data_ref(),
-        InputImpl->get_spec_const_blob_ref(), InputImpl->getOriginMask(),
-        std::move(RTCInfo), std::move(KernelNames));
+        InputImpl.get_kernel_ids_ptr(), Prog,
+        InputImpl.get_spec_const_data_ref(),
+        InputImpl.get_spec_const_blob_ref(), InputImpl.getOriginMask(),
+        std::move(RTCInfo), std::move(KernelNames),
+        /*MergedImageStorage = */ nullptr);
 
     std::string CompileOptions;
     applyCompileOptionsFromEnvironment(CompileOptions);
     appendCompileOptionsFromImage(
-        CompileOptions, *(InputImpl->get_bin_image_ref()), Devs, Adapter);
+        CompileOptions, *(InputImpl.get_bin_image_ref()), Devs, Adapter);
     // Should always come last!
     appendCompileEnvironmentVariablesThatAppend(CompileOptions);
-    ur_result_t Error = doCompile(
-        Adapter, ObjectImpl->get_ur_program_ref(), Devs.size(),
-        URDevices.data(),
-        getSyclObjImpl(InputImpl->get_context()).get()->getHandleRef(),
-        CompileOptions.c_str());
+    ur_result_t Error =
+        doCompile(Adapter, ObjectImpl->get_ur_program_ref(), Devs.size(),
+                  URDevices.data(),
+                  getSyclObjImpl(InputImpl.get_context()).get()->getHandleRef(),
+                  CompileOptions.c_str());
     if (Error != UR_RESULT_SUCCESS)
       throw sycl::exception(
           make_error_code(errc::build),
@@ -3074,13 +3075,11 @@ ProgramManager::link(const std::vector<device_image_plain> &Imgs,
   }
   auto MergedRTCInfo = detail::KernelCompilerBinaryInfo::Merge(RTCInfoPtrs);
 
-  DeviceImageImplPtr ExecutableImpl =
-      std::make_shared<detail::device_image_impl>(
-          NewBinImg, Context, std::vector<device>{Devs},
-          bundle_state::executable, std::move(KernelIDs), LinkedProg,
-          std::move(NewSpecConstMap), std::move(NewSpecConstBlob),
-          CombinedOrigins, std::move(MergedRTCInfo),
-          std::move(MergedKernelNames), std::move(MergedImageStorage));
+  DeviceImageImplPtr ExecutableImpl = device_image_impl::create(
+      NewBinImg, Context, std::vector<device>{Devs}, bundle_state::executable,
+      std::move(KernelIDs), LinkedProg, std::move(NewSpecConstMap),
+      std::move(NewSpecConstBlob), CombinedOrigins, std::move(MergedRTCInfo),
+      std::move(MergedKernelNames), std::move(MergedImageStorage));
 
   // TODO: Make multiple sets of device images organized by devices they are
   // compiled for.
@@ -3156,7 +3155,7 @@ ProgramManager::build(const DevImgPlainWithDeps &DevImgWithDeps,
   }
   auto MergedRTCInfo = detail::KernelCompilerBinaryInfo::Merge(RTCInfoPtrs);
 
-  DeviceImageImplPtr ExecImpl = std::make_shared<detail::device_image_impl>(
+  DeviceImageImplPtr ExecImpl = device_image_impl::create(
       ResultBinImg, Context, std::vector<device>{Devs},
       bundle_state::executable, std::move(KernelIDs), ResProgram,
       std::move(SpecConstMap), std::move(SpecConstBlob), CombinedOrigins,
diff --git a/sycl/source/detail/scheduler/commands.cpp b/sycl/source/detail/scheduler/commands.cpp
diff --git a/sycl/source/detail/scheduler/commands.hpp b/sycl/source/detail/scheduler/commands.hpp
