[NFC][SYCL] More use of device_images_range (#20692)

Author: aelovikov-intel

sycl/source/detail/device_image_impl.hpp

@@ -1361,6 +1361,21 @@ class device_images_range : public iterator_range<device_images_iterator> {
 
 public:
   using Base::Base;
+  template <typename Container>
+  decltype(std::declval<Base>().to<Container>()) to() const {
+    return this->Base::to<Container>();
+  }
+
+  template <typename Container>
+  std::enable_if_t<std::is_same_v<Container, std::vector<ur_program_handle_t>>,
+                   Container>
+  to() const {
+    std::vector<ur_program_handle_t> ProgramHandles;
+    ProgramHandles.reserve(size());
+    std::transform(begin(), end(), std::back_inserter(ProgramHandles),
+                   [](device_image_impl &Img) { return Img.get_ur_program(); });
+    return ProgramHandles;
+  }
 };
 
 } // namespace detail

sycl/source/detail/program_manager/program_manager.cpp

@@ -2908,24 +2908,22 @@ ProgramManager::compile(const DevImgPlainWithDeps &ImgWithDeps,
 // Returns a merged device binary image, new set of kernel IDs and new
 // specialization constant data.
 static const RTDeviceBinaryImage *
-mergeImageData(const std::vector<device_image_plain> &Imgs,
-               std::vector<kernel_id> &KernelIDs,
+mergeImageData(device_images_range Imgs, std::vector<kernel_id> &KernelIDs,
                std::vector<unsigned char> &NewSpecConstBlob,
                device_image_impl::SpecConstMapT &NewSpecConstMap,
                std::unique_ptr<DynRTDeviceBinaryImage> &MergedImageStorage) {
-  for (const device_image_plain &Img : Imgs) {
-    device_image_impl &DeviceImageImpl = *getSyclObjImpl(Img);
+  for (device_image_impl &Img : Imgs) {
     // Duplicates are not expected here, otherwise urProgramLink should fail
-    KernelIDs.insert(KernelIDs.end(), DeviceImageImpl.get_kernel_ids().begin(),
-                     DeviceImageImpl.get_kernel_ids().end());
+    KernelIDs.insert(KernelIDs.end(), Img.get_kernel_ids().begin(),
+                     Img.get_kernel_ids().end());
     // To be able to answer queries about specialziation constants, the new
     // device image should have the specialization constants from all the linked
     // images.
     const std::lock_guard<std::mutex> SpecConstLock(
-        DeviceImageImpl.get_spec_const_data_lock());
+        Img.get_spec_const_data_lock());
     // Copy all map entries to the new map. Since the blob will be copied to
     // the end of the new blob we need to move the blob offset of each entry.
-    for (const auto &SpecConstIt : DeviceImageImpl.get_spec_const_data_ref()) {
+    for (const auto &SpecConstIt : Img.get_spec_const_data_ref()) {
       std::vector<device_image_impl::SpecConstDescT> &NewDescEntries =
           NewSpecConstMap[SpecConstIt.first];
 
@@ -2943,21 +2941,21 @@ mergeImageData(const std::vector<device_image_plain> &Imgs,
     // Copy the blob from the device image into the new blob. This moves the
     // offsets of the following blobs.
     NewSpecConstBlob.insert(NewSpecConstBlob.end(),
-                            DeviceImageImpl.get_spec_const_blob_ref().begin(),
-                            DeviceImageImpl.get_spec_const_blob_ref().end());
+                            Img.get_spec_const_blob_ref().begin(),
+                            Img.get_spec_const_blob_ref().end());
   }
   // device_image_impl expects kernel ids to be sorted for fast search
   std::sort(KernelIDs.begin(), KernelIDs.end(), LessByHash<kernel_id>{});
 
   // If there is only a single image, use it as the result.
   if (Imgs.size() == 1)
-    return getSyclObjImpl(Imgs[0])->get_bin_image_ref();
+    return Imgs.front().get_bin_image_ref();
 
   // Otherwise we create a dynamic image with the merged information.
   std::vector<const RTDeviceBinaryImage *> BinImgs;
   BinImgs.reserve(Imgs.size());
-  for (const device_image_plain &Img : Imgs) {
-    auto ImgBinRef = getSyclObjImpl(Img)->get_bin_image_ref();
+  for (device_image_impl &Img : Imgs) {
+    auto ImgBinRef = Img.get_bin_image_ref();
     // For some cases, like SYCL kernel compiler binaries, we don't have
     // binaries. For these we assume no properties associated, so they can be
     // safely ignored.
@@ -2969,25 +2967,21 @@ mergeImageData(const std::vector<device_image_plain> &Imgs,
 }
 
 std::vector<device_image_plain>
-ProgramManager::link(const std::vector<device_image_plain> &Imgs,
-                     devices_range Devs, const property_list &PropList) {
+ProgramManager::link(device_images_range Imgs, devices_range Devs,
+                     const property_list &PropList) {
   {
     auto NoAllowedPropertiesCheck = [](int) { return false; };
     detail::PropertyValidator::checkPropsAndThrow(
         PropList, NoAllowedPropertiesCheck, NoAllowedPropertiesCheck);
   }
 
-  std::vector<ur_program_handle_t> URPrograms;
-  URPrograms.reserve(Imgs.size());
-  for (const device_image_plain &Img : Imgs)
-    URPrograms.push_back(getSyclObjImpl(Img)->get_ur_program());
-
+  auto URPrograms = Imgs.to<std::vector<ur_program_handle_t>>();
   auto URDevices = Devs.to<std::vector<ur_device_handle_t>>();
 
   // FIXME: Linker options are picked from the first object, but is that safe?
   std::string LinkOptionsStr;
   applyLinkOptionsFromEnvironment(LinkOptionsStr);
-  device_image_impl &FirstImgImpl = *getSyclObjImpl(Imgs[0]);
+  device_image_impl &FirstImgImpl = Imgs.front();
   if (LinkOptionsStr.empty() && FirstImgImpl.get_bin_image_ref())
     appendLinkOptionsFromImage(LinkOptionsStr,
                                *(FirstImgImpl.get_bin_image_ref()));
@@ -3045,8 +3039,8 @@ ProgramManager::link(const std::vector<device_image_plain> &Imgs,
     // underlying program disposed of). Protecting from incorrect values by
     // removal of map entries with same handle (obviously invalid entries).
     std::ignore = NativePrograms.erase(LinkedProg);
-    for (const device_image_plain &Img : Imgs) {
-      if (auto BinImageRef = getSyclObjImpl(Img)->get_bin_image_ref())
+    for (device_image_impl &Img : Imgs) {
+      if (auto BinImageRef = Img.get_bin_image_ref())
         NativePrograms.insert(
             {LinkedProg, {ContextImpl.shared_from_this(), BinImageRef}});
     }
@@ -3062,15 +3056,14 @@ ProgramManager::link(const std::vector<device_image_plain> &Imgs,
   KernelNameSetT MergedKernelNames;
   std::map<std::string, KernelArgMask, std::less<>>
       MergedEliminatedKernelArgMasks;
-  for (const device_image_plain &DevImg : Imgs) {
-    device_image_impl &DevImgImpl = *getSyclObjImpl(DevImg);
-    CombinedOrigins |= DevImgImpl.getOriginMask();
-    RTCInfoPtrs.emplace_back(&(DevImgImpl.getRTCInfo()));
-    MergedKernelNames.insert(DevImgImpl.getKernelNames().begin(),
-                             DevImgImpl.getKernelNames().end());
+  for (device_image_impl &DevImg : Imgs) {
+    CombinedOrigins |= DevImg.getOriginMask();
+    RTCInfoPtrs.emplace_back(&(DevImg.getRTCInfo()));
+    MergedKernelNames.insert(DevImg.getKernelNames().begin(),
+                             DevImg.getKernelNames().end());
     MergedEliminatedKernelArgMasks.insert(
-        DevImgImpl.getEliminatedKernelArgMasks().begin(),
-        DevImgImpl.getEliminatedKernelArgMasks().end());
+        DevImg.getEliminatedKernelArgMasks().begin(),
+        DevImg.getEliminatedKernelArgMasks().end());
   }
   auto MergedRTCInfo = detail::KernelCompilerBinaryInfo::Merge(RTCInfoPtrs);
 

sycl/source/detail/program_manager/program_manager.hpp

@@ -76,6 +76,7 @@ class device_impl;
 class devices_range;
 class queue_impl;
 class event_impl;
+class device_images_range;
 // DeviceLibExt is shared between sycl runtime and sycl-post-link tool.
 // If any update is made here, need to sync with DeviceLibExt definition
 // in llvm/tools/sycl-post-link/sycl-post-link.cpp
@@ -355,9 +356,9 @@ class ProgramManager {
 
   // Produces set of device images by convering input device images to object
   // the executable state
-  std::vector<device_image_plain>
-  link(const std::vector<device_image_plain> &Imgs, devices_range Devs,
-       const property_list &PropList);
+  std::vector<device_image_plain> link(device_images_range Imgs,
+                                       devices_range Devs,
+                                       const property_list &PropList);
 
   // Produces new device image by converting input device image to the
   // executable state