Delay image decompression till it is actually used.

Author: uditagarwal97

clang/tools/clang-offload-wrapper/ClangOffloadWrapper.cpp

@@ -1112,10 +1112,11 @@ class BinaryWrapper {
         // If '--offload-compress' option is specified and zstd is not
         // available, throw an error.
         if (OffloadCompressDevImgs && !llvm::compression::zstd::isAvailable()) {
-          createStringError(inconvertibleErrorCode(),
-                            "'--offload-compress' option is specified but zstd "
-                            "is not available. The device image will not be "
-                            "compressed.");
+          return createStringError(
+              inconvertibleErrorCode(),
+              "'--offload-compress' option is specified but zstd "
+              "is not available. The device image will not be "
+              "compressed.");
         }
 
         // Don't compress if the user explicitly specifies the binary image

sycl/source/detail/compression.hpp

@@ -81,6 +81,17 @@ class ZSTDCompressor {
     return dstBuffer;
   }
 
+  static size_t GetDecompressedSize(const char *src, size_t srcSize) {
+    size_t dstBufferSize = ZSTD_getFrameContentSize(src, srcSize);
+
+    if (dstBufferSize == ZSTD_CONTENTSIZE_UNKNOWN ||
+        dstBufferSize == ZSTD_CONTENTSIZE_ERROR) {
+      throw sycl::exception(sycl::make_error_code(sycl::errc::runtime),
+                            "Error determining size of uncompressed data.");
+    }
+    return dstBufferSize;
+  }
+
   static std::unique_ptr<char> DecompressBlob(const char *src, size_t srcSize,
                                               size_t &dstSize) {
     auto &instance = GetSingletonInstance();
@@ -101,13 +112,7 @@ class ZSTDCompressor {
 
     // Size of decompressed image can be larger than what we can allocate
     // on heap. In that case, we need to use streaming decompression.
-    auto dstBufferSize = ZSTD_getFrameContentSize(src, srcSize);
-
-    if (dstBufferSize == ZSTD_CONTENTSIZE_UNKNOWN ||
-        dstBufferSize == ZSTD_CONTENTSIZE_ERROR) {
-      throw sycl::exception(sycl::make_error_code(sycl::errc::runtime),
-                            "Error determining size of uncompressed data.");
-    }
+    auto dstBufferSize = GetDecompressedSize(src, srcSize);
 
     // Allocate buffer for decompressed data.
     auto dstBuffer = std::unique_ptr<char>(new char[dstBufferSize]);

sycl/source/detail/device_binary_image.cpp

@@ -170,6 +170,8 @@ void RTDeviceBinaryImage::init(sycl_device_binary Bin) {
   // it when invoking the offload wrapper job
   Format = static_cast<ur::DeviceBinaryType>(Bin->Format);
 
+  // For compressed images, we delay determining the format until the image is
+  // decompressed.
   if (Format == SYCL_DEVICE_BINARY_TYPE_NONE)
     // try to determine the format; may remain "NONE"
     Format = ur::getBinaryImageFormat(Bin->BinaryStart, getSize());
@@ -186,7 +188,6 @@ void RTDeviceBinaryImage::init(sycl_device_binary Bin) {
     ProgramMetadataUR.push_back(
         ur::mapDeviceBinaryPropertyToProgramMetadata(Prop));
   }
-
   ExportedSymbols.init(Bin, __SYCL_PROPERTY_SET_SYCL_EXPORTED_SYMBOLS);
   ImportedSymbols.init(Bin, __SYCL_PROPERTY_SET_SYCL_IMPORTED_SYMBOLS);
   DeviceGlobals.init(Bin, __SYCL_PROPERTY_SET_SYCL_DEVICE_GLOBALS);
@@ -235,25 +236,34 @@ CompressedRTDeviceBinaryImage::CompressedRTDeviceBinaryImage(
     sycl_device_binary CompressedBin)
     : RTDeviceBinaryImage() {
 
-  size_t compressedDataSize = static_cast<size_t>(CompressedBin->BinaryEnd -
-                                                  CompressedBin->BinaryStart);
+  // 'CompressedBin' is part of the executable image loaded into memory
+  // which can't be modified easily. So, we need to make a copy of it.
+  Bin = new sycl_device_binary_struct(*CompressedBin);
+
+  // Get the decompressed size of the binary image.
+  m_ImageSize = ZSTDCompressor::GetDecompressedSize(
+      reinterpret_cast<const char *>(Bin->BinaryStart),
+      static_cast<size_t>(Bin->BinaryEnd - Bin->BinaryStart));
+
+  init(Bin);
+}
+
+void CompressedRTDeviceBinaryImage::Decompress() {
+
+  size_t CompressedDataSize =
+      static_cast<size_t>(Bin->BinaryEnd - Bin->BinaryStart);
 
   size_t DecompressedSize = 0;
   m_DecompressedData = ZSTDCompressor::DecompressBlob(
-      reinterpret_cast<const char *>(CompressedBin->BinaryStart),
-      compressedDataSize, DecompressedSize);
+      reinterpret_cast<const char *>(Bin->BinaryStart), CompressedDataSize,
+      DecompressedSize);
 
-  Bin = new sycl_device_binary_struct(*CompressedBin);
   Bin->BinaryStart =
       reinterpret_cast<const unsigned char *>(m_DecompressedData.get());
   Bin->BinaryEnd = Bin->BinaryStart + DecompressedSize;
 
-  // Set the new format to none and let RT determine the format.
-  // TODO: Add support for automatically detecting compressed
-  // binary format.
-  Bin->Format = SYCL_DEVICE_BINARY_TYPE_NONE;
-
-  init(Bin);
+  Bin->Format = ur::getBinaryImageFormat(Bin->BinaryStart, getSize());
+  Format = static_cast<ur::DeviceBinaryType>(Bin->Format);
 }
 
 CompressedRTDeviceBinaryImage::~CompressedRTDeviceBinaryImage() {

sycl/source/detail/device_binary_image.hpp

@@ -158,7 +158,10 @@ class RTDeviceBinaryImage {
   virtual void print() const;
   virtual void dump(std::ostream &Out) const;
 
-  size_t getSize() const {
+  // getSize will be overridden in the case of compressed binary images.
+  // In that case, we return the size of uncompressed data, instead of
+  // BinaryEnd - BinaryStart.
+  virtual size_t getSize() const {
     assert(Bin && "binary image data not set");
     return static_cast<size_t>(Bin->BinaryEnd - Bin->BinaryStart);
   }
@@ -277,21 +280,31 @@ class DynRTDeviceBinaryImage : public RTDeviceBinaryImage {
 };
 
 #ifndef SYCL_RT_ZSTD_NOT_AVAIABLE
-// Compressed device binary image. It decompresses the binary image on
-// construction and stores the decompressed data as RTDeviceBinaryImage.
+// Compressed device binary image. Decompression happens when the image is
+// actually used to build a program.
 // Also, frees the decompressed data in destructor.
 class CompressedRTDeviceBinaryImage : public RTDeviceBinaryImage {
 public:
   CompressedRTDeviceBinaryImage(sycl_device_binary Bin);
   ~CompressedRTDeviceBinaryImage() override;
 
+  void Decompress();
+
+  // We return the size of decompressed data, not the size of compressed data.
+  size_t getSize() const override {
+    assert(Bin && "binary image data not set");
+    return m_ImageSize;
+  }
+
+  bool IsCompressed() const { return m_DecompressedData.get() == nullptr; }
   void print() const override {
     RTDeviceBinaryImage::print();
     std::cerr << "    COMPRESSED\n";
   }
 
 private:
   std::unique_ptr<char> m_DecompressedData;
+  size_t m_ImageSize;
 };
 #endif // SYCL_RT_ZSTD_NOT_AVAIABLE
 

sycl/source/detail/program_manager/program_manager.cpp

@@ -733,6 +733,12 @@ setSpecializationConstants(const std::shared_ptr<device_image_impl> &InputImpl,
   }
 }
 
+static inline void CheckAndDecompressImage(RTDeviceBinaryImage *Img) {
+  if (auto CompImg = dynamic_cast<CompressedRTDeviceBinaryImage *>(Img))
+    if (CompImg->IsCompressed())
+      CompImg->Decompress();
+}
+
 // When caching is enabled, the returned UrProgram will already have
 // its ref count incremented.
 ur_program_handle_t ProgramManager::getBuiltURProgram(
@@ -785,6 +791,10 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
       collectDeviceImageDepsForImportedSymbols(Img, Device);
   DeviceImagesToLink.insert(ImageDeps.begin(), ImageDeps.end());
 
+  // Decompress all DeviceImagesToLink
+  for (RTDeviceBinaryImage *BinImg : DeviceImagesToLink)
+    CheckAndDecompressImage(BinImg);
+
   std::vector<const RTDeviceBinaryImage *> AllImages;
   AllImages.reserve(ImageDeps.size() + 1);
   AllImages.push_back(&Img);
@@ -1388,6 +1398,10 @@ ProgramManager::getDeviceImage(const std::string &KernelName,
                                     Device);
     }
   }
+
+  // Decompress the image if it is compressed.
+  CheckAndDecompressImage(Img);
+
   if (Img) {
     CheckJITCompilationForImage(Img, JITCompilationIsRequired);
 
@@ -1714,6 +1728,10 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
           [&](auto &CurrentImg) {
             return CurrentImg.first->getFormat() == Img->getFormat();
           });
+
+      // Check if image is compressed, and decompress it before dumping.
+      CheckAndDecompressImage(Img.get());
+
       dumpImage(*Img, NeedsSequenceID ? ++SequenceID : 0);
     }
 
@@ -2191,6 +2209,9 @@ ProgramManager::getSYCLDeviceImagesWithCompatibleState(
 
         auto &[KernelImagesState, KernelImages] = *StateImagesPair;
 
+        // Check if device image is compressed and decompress it if needed
+        CheckAndDecompressImage(BinImage);
+
         if (KernelImages.empty()) {
           KernelImagesState = ImgState;
           KernelImages.push_back(BinImage);

sycl/test-e2e/Compression/compression_multiple_tu.cpp

@@ -3,10 +3,10 @@
 // REQUIRES: zstd
 
 // DEFINE: %{fPIC_flag} =  %if windows %{%} %else %{-fPIC%}
-// RUN: %{build} --offload-compress -DENABLE_KERNEL1 -shared %{fPIC_flag} -o %t_kernel1.so
-// RUN: %{build} -DENABLE_KERNEL2 -shared %{fPIC_flag} -o %t_kernel2.so
+// RUN: %{build} --offload-compress -DENABLE_KERNEL1 -shared %{fPIC_flag} -o %T/kernel1.so
+// RUN: %{build} -DENABLE_KERNEL2 -shared %{fPIC_flag} -o %T/kernel2.so
 
-// RUN: %{build} %t_kernel1.so %t_kernel2.so -Wl,-rpath=%T -o %t_compress.out
+// RUN: %{build} %t_kernel1.so %t_kernel2.so -o %t_compress.out
 // RUN: %{run} %t_compress.out
 #if defined(ENABLE_KERNEL1) || defined(ENABLE_KERNEL2)
 #include <sycl/builtins.hpp>