[SYCL] Fix linking of compressed device images when dependencies are not compressed

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