diff --git a/sycl/source/detail/program_manager/program_manager.cpp b/sycl/source/detail/program_manager/program_manager.cpp index 885c019cf7050..d011132f1c5df 100644 --- a/sycl/source/detail/program_manager/program_manager.cpp +++ b/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(Img)) + if (CompImg->IsCompressed()) + CompImg->Decompress(); +#endif +} + std::set 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(&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 &InputImpl, } } -static inline void -CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) { -#ifndef SYCL_RT_ZSTD_NOT_AVAIABLE - if (auto CompImg = dynamic_cast(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( diff --git a/sycl/test-e2e/DeviceImageDependencies/dynamic_compress.cpp b/sycl/test-e2e/DeviceImageDependencies/dynamic_compress.cpp new file mode 100644 index 0000000000000..9cc9c86710f4a --- /dev/null +++ b/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