[SYCL] Fix linking of compressed device images when dependencies are not compressed (#18906)

**Problem**
When linking device images, we reject dependencies whose image format
does not match the parent image. However, consider the case when parent
image is compressed, while dependencies are not (demonstrated in the
test case attached to this PR). In this case, we are incorrectly
rejecting device images and thus causing `No device image found for
external symbol` error.

**Solution**
If the format of the main and dependent device image differs and one of
them is compressed, we decompress them and recheck the format of
decompressed device images.
One side-effect of this solution is that now we'll have to decompress
device images, even if we are not using them. For example, when format
of decompressed main and dependent images differs. Unfortunately,
there's no way to find format of the compressed device image, without
first decompressing it.
However, I don't think this will incur a significant overhead as (1) we
decompress device image only once and cache it for subsequent use, and
(2) we decompress only if the dependent device image has an export
symbol that main device image wants (when finding which images to link)
and if it is compatible with the device.

Author: uditagarwal97

sycl/source/detail/program_manager/program_manager.cpp

@@ -687,6 +687,15 @@ ProgramManager::collectDeviceImageDeps(const RTDeviceBinaryImage &Img,
   return DeviceImagesToLink;
 }
 
+static inline void
+CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) {
+#ifndef SYCL_RT_ZSTD_NOT_AVAIABLE
+  if (auto CompImg = dynamic_cast<CompressedRTDeviceBinaryImage *>(Img))
+    if (CompImg->IsCompressed())
+      CompImg->Decompress();
+#endif
+}
+
 std::set<RTDeviceBinaryImage *>
 ProgramManager::collectDeviceImageDepsForImportedSymbols(
     const RTDeviceBinaryImage &MainImg, const device &Dev,
@@ -713,13 +722,30 @@ ProgramManager::collectDeviceImageDepsForImportedSymbols(
     bool Found = false;
     for (auto It = Range.first; It != Range.second; ++It) {
       RTDeviceBinaryImage *Img = It->second;
-      if (Img->getFormat() != Format ||
-          !doesDevSupportDeviceRequirements(Dev, *Img) ||
+
+      if (!doesDevSupportDeviceRequirements(Dev, *Img) ||
           !compatibleWithDevice(Img, *getSyclObjImpl(Dev).get()))
         continue;
+
+      // If the image is a special device image, we need to check if it
+      // should be used for this device.
       if (isSpecialDeviceImage(Img) &&
           !isSpecialDeviceImageShouldBeUsed(Img, *getSyclObjImpl(Dev).get()))
         continue;
+
+      // If any of the images is compressed, we need to decompress it
+      // and then check if the format matches.
+      if (Format == SYCL_DEVICE_BINARY_TYPE_COMPRESSED_NONE ||
+          Img->getFormat() == SYCL_DEVICE_BINARY_TYPE_COMPRESSED_NONE) {
+        auto MainImgPtr = const_cast<RTDeviceBinaryImage *>(&MainImg);
+        CheckAndDecompressImage(MainImgPtr);
+        CheckAndDecompressImage(Img);
+        Format = MainImg.getFormat();
+      }
+      // Skip this image if its format differs from the main image.
+      if (Img->getFormat() != Format)
+        continue;
+
       DeviceImagesToLink.insert(Img);
       Found = true;
       for (const sycl_device_binary_property &ISProp :
@@ -836,15 +862,6 @@ setSpecializationConstants(const std::shared_ptr<device_image_impl> &InputImpl,
   }
 }
 
-static inline void
-CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) {
-#ifndef SYCL_RT_ZSTD_NOT_AVAIABLE
-  if (auto CompImg = dynamic_cast<CompressedRTDeviceBinaryImage *>(Img))
-    if (CompImg->IsCompressed())
-      CompImg->Decompress();
-#endif
-}
-
 // When caching is enabled, the returned UrProgram will already have
 // its ref count incremented.
 ur_program_handle_t ProgramManager::getBuiltURProgram(

sycl/test-e2e/DeviceImageDependencies/dynamic_compress.cpp

@@ -0,0 +1,25 @@
+// Test device image linking when using dynamic libraries when one of
+// the device image is compressed and the other is not.
+
+// UNSUPPORTED: target-nvidia || target-amd
+// UNSUPPORTED-INTENDED: Linking using dynamic libraries is not supported on AMD
+// and Nvidia.
+// REQUIRES: zstd
+
+// DEFINE: %{dynamic_lib_options} = -fsycl %fPIC %shared_lib -fsycl-allow-device-image-dependencies -I %S/Inputs %if windows %{-DMAKE_DLL %}
+// DEFINE: %{dynamic_lib_suffix} = %if windows %{dll%} %else %{so%}
+
+// RUN: %clangxx %{dynamic_lib_options} %S/Inputs/d.cpp                                    -o %T/libdevice_d.%{dynamic_lib_suffix}
+// RUN: %clangxx %{dynamic_lib_options} %S/Inputs/c.cpp %if windows %{%T/libdevice_d.lib%} -o %T/libdevice_c.%{dynamic_lib_suffix}
+// RUN: %clangxx %{dynamic_lib_options} %S/Inputs/b.cpp %if windows %{%T/libdevice_c.lib%} -o %T/libdevice_b.%{dynamic_lib_suffix}
+// RUN: %clangxx %{dynamic_lib_options} %S/Inputs/a.cpp %if windows %{%T/libdevice_b.lib%} -o %T/libdevice_a.%{dynamic_lib_suffix}
+
+// Compressed main executable, while dependencies are not compressed.
+
+// RUN: %clangxx -fsycl --offload-compress %{sycl_target_opts} -fsycl-allow-device-image-dependencies -fsycl-device-code-split=per_kernel %S/Inputs/basic.cpp -o %t.out            \
+// RUN: %if windows                                                                       \
+// RUN:   %{%T/libdevice_a.lib%}                                                          \
+// RUN: %else                                                                             \
+// RUN:   %{-L%T -ldevice_a -ldevice_b -ldevice_c -ldevice_d -Wl,-rpath=%T%}
+
+// RUN: %{run} %t.out