[NFC][SYCL] Prepare kernel_bundle_impl.hpp for getSyclObjImpl to return raw ref (#19251)

I'm planning to change `getSyclObjImpl` to return a raw reference in a
later patch, uploading a bunch of PRs in preparation to that to make the
subsequent review easier.

Author: aelovikov-intel

sycl/source/detail/kernel_bundle_impl.hpp

@@ -138,12 +138,11 @@ class kernel_bundle_impl
       : MContext(InputBundle.get_context()), MDevices(std::move(Devs)),
         MState(TargetState) {
 
-    const std::shared_ptr<kernel_bundle_impl> &InputBundleImpl =
-        getSyclObjImpl(InputBundle);
-    MSpecConstValues = InputBundleImpl->get_spec_const_map_ref();
+    kernel_bundle_impl &InputBundleImpl = *getSyclObjImpl(InputBundle);
+    MSpecConstValues = InputBundleImpl.get_spec_const_map_ref();
 
     const std::vector<device> &InputBundleDevices =
-        InputBundleImpl->get_devices();
+        InputBundleImpl.get_devices();
     const bool AllDevsAssociatedWithInputBundle =
         std::all_of(MDevices.begin(), MDevices.end(),
                     [&InputBundleDevices](const device &Dev) {
@@ -158,11 +157,11 @@ class kernel_bundle_impl
           "devices for input bundle or vector of devices is empty");
 
     // Copy SYCLBINs to ensure lifetime is preserved by the executable bundle.
-    MSYCLBINs.insert(MSYCLBINs.end(), InputBundleImpl->MSYCLBINs.begin(),
-                     InputBundleImpl->MSYCLBINs.end());
+    MSYCLBINs.insert(MSYCLBINs.end(), InputBundleImpl.MSYCLBINs.begin(),
+                     InputBundleImpl.MSYCLBINs.end());
 
     for (const DevImgPlainWithDeps &DevImgWithDeps :
-         InputBundleImpl->MDeviceImages) {
+         InputBundleImpl.MDeviceImages) {
       // Skip images which are not compatible with devices provided
       if (std::none_of(MDevices.begin(), MDevices.end(),
                        [&DevImgWithDeps](const device &Dev) {
@@ -311,11 +310,11 @@ class kernel_bundle_impl
     // images collection.
     std::map<std::string_view, size_t> ExportMap;
     for (size_t I = 0; I < DevImages.size(); ++I) {
-      auto DevImageImpl = getSyclObjImpl(DevImages[I]);
-      if (DevImageImpl->get_bin_image_ref() == nullptr)
+      device_image_impl &DevImageImpl = *getSyclObjImpl(DevImages[I]);
+      if (DevImageImpl.get_bin_image_ref() == nullptr)
         continue;
       for (const sycl_device_binary_property &ESProp :
-           DevImageImpl->get_bin_image_ref()->getExportedSymbols()) {
+           DevImageImpl.get_bin_image_ref()->getExportedSymbols()) {
         if (ExportMap.find(ESProp->Name) != ExportMap.end())
           throw sycl::exception(make_error_code(errc::invalid),
                                 "Duplicate exported symbol \"" +
@@ -329,12 +328,12 @@ class kernel_bundle_impl
     std::vector<std::vector<size_t>> Dependencies;
     Dependencies.resize(DevImages.size());
     for (size_t I = 0; I < DevImages.size(); ++I) {
-      auto DevImageImpl = getSyclObjImpl(DevImages[I]);
-      if (DevImageImpl->get_bin_image_ref() == nullptr)
+      device_image_impl &DevImageImpl = *getSyclObjImpl(DevImages[I]);
+      if (DevImageImpl.get_bin_image_ref() == nullptr)
         continue;
       std::set<size_t> DeviceImageDepsSet;
       for (const sycl_device_binary_property &ISProp :
-           DevImageImpl->get_bin_image_ref()->getImportedSymbols()) {
+           DevImageImpl.get_bin_image_ref()->getImportedSymbols()) {
         auto ExportSymbolIt = ExportMap.find(ISProp->Name);
         if (ExportSymbolIt == ExportMap.end())
           throw sycl::exception(make_error_code(errc::invalid),
@@ -348,13 +347,12 @@ class kernel_bundle_impl
     }
 
     // Create a link graph and clone it for each device.
-    const std::shared_ptr<device_impl> &FirstDevice =
-        getSyclObjImpl(MDevices[0]);
+    device_impl &FirstDevice = *getSyclObjImpl(MDevices[0]);
     std::map<std::shared_ptr<device_impl>, LinkGraph<device_image_plain>>
         DevImageLinkGraphs;
     const auto &FirstGraph =
         DevImageLinkGraphs
-            .emplace(FirstDevice,
+            .emplace(FirstDevice.shared_from_this(),
                      LinkGraph<device_image_plain>{DevImages, Dependencies})
             .first->second;
     for (size_t I = 1; I < MDevices.size(); ++I)
@@ -498,12 +496,12 @@ class kernel_bundle_impl
     if (get_bundle_state() == bundle_state::input) {
       // Copy spec constants values from the device images.
       auto MergeSpecConstants = [this](const device_image_plain &Img) {
-        const detail::DeviceImageImplPtr &ImgImpl = getSyclObjImpl(Img);
+        detail::device_image_impl &ImgImpl = *getSyclObjImpl(Img);
         const std::map<std::string,
                        std::vector<device_image_impl::SpecConstDescT>>
-            &SpecConsts = ImgImpl->get_spec_const_data_ref();
+            &SpecConsts = ImgImpl.get_spec_const_data_ref();
         const std::vector<unsigned char> &Blob =
-            ImgImpl->get_spec_const_blob_ref();
+            ImgImpl.get_spec_const_blob_ref();
         for (const std::pair<const std::string,
                              std::vector<device_image_impl::SpecConstDescT>>
                  &SpecConst : SpecConsts) {
@@ -675,10 +673,9 @@ class kernel_bundle_impl
     //       resulting kernel object should be able to map devices to their
     //       respective backend kernel objects.
     for (const device_image_plain &DevImg : MUniqueDeviceImages) {
-      const std::shared_ptr<device_image_impl> &DevImgImpl =
-          getSyclObjImpl(DevImg);
+      device_image_impl &DevImgImpl = *getSyclObjImpl(DevImg);
       if (std::shared_ptr<kernel_impl> PotentialKernelImpl =
-              DevImgImpl->tryGetExtensionKernel(Name, MContext, *this))
+              DevImgImpl.tryGetExtensionKernel(Name, MContext, *this))
         return detail::createSyclObjFromImpl<kernel>(
             std::move(PotentialKernelImpl));
     }
@@ -731,10 +728,10 @@ class kernel_bundle_impl
                             "'device_image_scope' property");
     }
 
-    const auto &DeviceImpl = getSyclObjImpl(Dev);
+    device_impl &DeviceImpl = *getSyclObjImpl(Dev);
     bool SupportContextMemcpy = false;
-    DeviceImpl->getAdapter()->call<UrApiKind::urDeviceGetInfo>(
-        DeviceImpl->getHandleRef(),
+    DeviceImpl.getAdapter()->call<UrApiKind::urDeviceGetInfo>(
+        DeviceImpl.getHandleRef(),
         UR_DEVICE_INFO_USM_CONTEXT_MEMCPY_SUPPORT_EXP,
         sizeof(SupportContextMemcpy), &SupportContextMemcpy, nullptr);
     if (SupportContextMemcpy) {
@@ -764,14 +761,14 @@ class kernel_bundle_impl
     // Collect kernel ids from all device images, then remove duplicates
     std::vector<kernel_id> Result;
     for (const device_image_plain &DeviceImage : MUniqueDeviceImages) {
-      const auto &DevImgImpl = getSyclObjImpl(DeviceImage);
+      detail::device_image_impl &DevImgImpl = *getSyclObjImpl(DeviceImage);
 
       // RTC kernel bundles shouldn't have user-facing kernel ids, return an
       // empty vector when the bundle contains RTC kernels.
-      if (DevImgImpl->getRTCInfo())
+      if (DevImgImpl.getRTCInfo())
         continue;
 
-      const std::vector<kernel_id> &KernelIDs = DevImgImpl->get_kernel_ids();
+      const std::vector<kernel_id> &KernelIDs = DevImgImpl.get_kernel_ids();
 
       Result.insert(Result.end(), KernelIDs.begin(), KernelIDs.end());
     }
@@ -1016,10 +1013,9 @@ class kernel_bundle_impl
     //       {kernel_name, device} and their corresponding image.
     // First look through the kernels registered in source-based images.
     for (const device_image_plain &DevImg : MUniqueDeviceImages) {
-      const std::shared_ptr<device_image_impl> &DevImgImpl =
-          getSyclObjImpl(DevImg);
+      device_image_impl &DevImgImpl = *getSyclObjImpl(DevImg);
       if (std::shared_ptr<kernel_impl> SourceBasedKernel =
-              DevImgImpl->tryGetExtensionKernel(Name, MContext, *this))
+              DevImgImpl.tryGetExtensionKernel(Name, MContext, *this))
         return SourceBasedKernel;
     }