[SYCL][Clang] Add support for device image compression

buildbot/configure.py

@@ -178,6 +178,7 @@ def do_configure(args):
         "-DLLVM_ENABLE_PROJECTS={}".format(llvm_enable_projects),
         "-DSYCL_BUILD_PI_HIP_PLATFORM={}".format(sycl_build_pi_hip_platform),
         "-DLLVM_BUILD_TOOLS=ON",
+        "-DLLVM_ENABLE_ZSTD=ON",
         "-DSYCL_ENABLE_WERROR={}".format(sycl_werror),
         "-DCMAKE_INSTALL_PREFIX={}".format(install_dir),
         "-DSYCL_INCLUDE_TESTS=ON",  # Explicitly include all kinds of SYCL tests.

clang/lib/Driver/ToolChains/Clang.cpp

@@ -10087,6 +10087,19 @@ void OffloadWrapper::ConstructJob(Compilation &C, const JobAction &JA,
     SmallString<128> TargetTripleOpt = TT.getArchName();
     bool WrapFPGADevice = false;
     bool FPGAEarly = false;
+
+    // Validate and propogate CLI options related to device image compression.
+    // -offload-compress
+    if (C.getInputArgs().getLastArg(options::OPT_offload_compress)) {
+      WrapperArgs.push_back(
+          C.getArgs().MakeArgString(Twine("-offload-compress")));
+      // -offload-compression-level=<>
+      if (Arg *A = C.getInputArgs().getLastArg(
+              options::OPT_offload_compression_level_EQ))
+        WrapperArgs.push_back(C.getArgs().MakeArgString(
+            Twine("-offload-compression-level=") + A->getValue()));
+    }
+
     if (Arg *A = C.getInputArgs().getLastArg(options::OPT_fsycl_link_EQ)) {
       WrapFPGADevice = true;
       FPGAEarly = (A->getValue() == StringRef("early"));

clang/test/Driver/clang-offload-wrapper-zstd.c

@@ -0,0 +1,40 @@
+// REQUIRES: zstd && (system-windows || system-linux)
+
+// clang-offload-wrapper compression test: checks that the wrapper can compress the device images.
+// Checks the '--offload-compress', '--offload-compression-level', and '--offload-compression-threshold'
+// CLI options.
+
+// --- Prepare test data by creating the debice binary image.
+// RUN: echo -e -n 'device binary image1\n' > %t.bin
+// RUN: echo -e -n '[Category1]\nint_prop1=1|10\n[Category2]\nint_prop2=1|20\n' > %t.props
+// RUN: echo -e -n 'kernel1\nkernel2\n' > %t.sym
+// RUN: echo -e -n 'Manifest file - arbitrary data generated by the toolchain\n' > %t.mnf
+// RUN: echo '[Code|Properties|Symbols|Manifest]' > %t.img1
+// RUN: echo %t.bin"|"%t.props"|"%t.sym"|"%t.mnf >> %t.img1
+
+///////////////////////////////////////////////////////
+// Compress the test image using clang-offload-wrapper.
+///////////////////////////////////////////////////////
+
+// RUN: clang-offload-wrapper -kind=sycl -target=TARGET -batch %t.img1 -o %t.wrapped.bc -v \
+// RUN: --offload-compress --offload-compression-level=9 --offload-compression-threshold=0 \
+// RUN: 2>&1 | FileCheck %s --check-prefix=CHECK-COMPRESS
+
+// CHECK-COMPRESS: [Compression] Original image size:
+// CHECK-COMPRESS: [Compression] Compressed image size:
+// CHECK-COMPRESS: [Compression] Compression level used: 9
+
+///////////////////////////////////////////////////////////
+// Check that there is no compression when the threshold is set to a value higher than the image size
+// or '--offload-compress' is not set.
+///////////////////////////////////////////////////////////
+
+// RUN: clang-offload-wrapper -kind=sycl -target=TARGET -batch %t.img1 -o %t.wrapped.bc -v \
+// RUN: --offload-compress --offload-compression-level=3 --offload-compression-threshold=1000 \
+// RUN: 2>&1 | FileCheck %s --check-prefix=CHECK-NO-COMPRESS
+
+// RUN: clang-offload-wrapper -kind=sycl -target=TARGET -batch %t.img1 -o %t.wrapped.bc -v \
+// RUN: --offload-compression-level=3 --offload-compression-threshold=0 \
+// RUN: 2>&1 | FileCheck %s --check-prefix=CHECK-NO-COMPRESS
+
+// CHECK-NO-COMPRESS-NOT: [Compression] Original image size:

clang/tools/clang-offload-wrapper/CMakeLists.txt

@@ -10,6 +10,7 @@ add_clang_tool(clang-offload-wrapper
 
 set(CLANG_OFFLOAD_WRAPPER_LIB_DEPS
   clangBasic
+  LLVMSupport
   )
 
 add_dependencies(clang clang-offload-wrapper)

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

@@ -67,6 +67,9 @@
 #include <string>
 #include <tuple>
 
+// For device image compression.
+#include <llvm/Support/Compression.h>
+
 #define OPENMP_OFFLOAD_IMAGE_VERSION "1.0"
 
 using namespace llvm;
@@ -139,15 +142,35 @@ static cl::list<std::string> Inputs(cl::Positional, cl::OneOrMore,
                                     cl::desc("<input files>"),
                                     cl::cat(ClangOffloadWrapperCategory));
 
+// CLI options for device image compression.
+static cl::opt<bool> OffloadCompressDevImgs(
+    "offload-compress", cl::init(false), cl::Optional,
+    cl::desc("Enable device image compression using ZSTD."),
+    cl::cat(ClangOffloadWrapperCategory));
+
+static cl::opt<int>
+    OffloadCompressLevel("offload-compression-level", cl::init(10),
+                         cl::Optional,
+                         cl::desc("ZSTD Compression level. Default: 10"),
+                         cl::cat(ClangOffloadWrapperCategory));
+
+static cl::opt<int>
+    OffloadCompressThreshold("offload-compression-threshold", cl::init(512),
+                             cl::Optional,
+                             cl::desc("Threshold (in bytes) over which to "
+                                      "compress images. Default: 512"),
+                             cl::cat(ClangOffloadWrapperCategory));
+
 // Binary image formats supported by this tool. The support basically means
 // mapping string representation given at the command line to a value from this
 // enum. No format checking is performed.
 enum BinaryImageFormat {
   none,   // image kind is not determined
   native, // image kind is native
   // portable image kinds go next
-  spirv, // SPIR-V
-  llvmbc // LLVM bitcode
+  spirv,          // SPIR-V
+  llvmbc,         // LLVM bitcode
+  compressed_none // compressed image with unknown format
 };
 
 /// Sets offload kind.
@@ -265,6 +288,8 @@ static StringRef formatToString(BinaryImageFormat Fmt) {
     return "llvmbc";
   case BinaryImageFormat::native:
     return "native";
+  case BinaryImageFormat::compressed_none:
+    return "compressed_none";
   }
   llvm_unreachable("bad format");
 
@@ -1083,10 +1108,55 @@ class BinaryWrapper {
           return FBinOrErr.takeError();
         Fbin = *FBinOrErr;
       } else {
-        Fbin = addDeviceImageToModule(
-            ArrayRef<char>(Bin->getBufferStart(), Bin->getBufferSize()),
-            Twine(OffloadKindTag) + Twine(ImgId) + Twine(".data"), Kind,
-            Img.Tgt);
+
+        // If '--offload-compress' option is specified and zstd is not available
+        // then warn the user that the image will not be compressed.
+        if (OffloadCompressDevImgs && !llvm::compression::zstd::isAvailable()) {
+          WithColor::warning(errs(), ToolName)
+              << "'--offload-compress' option is specified but zstd is not "
+                 "available. The device image will not be compressed.\n";
+        }
+
+        // Don't compress if the user explicitly specifies the binary image
+        // format or if the image is smaller than OffloadCompressThreshold
+        // bytes.
+        if (Kind != OffloadKind::SYCL || !OffloadCompressDevImgs ||
+            Img.Fmt != BinaryImageFormat::none ||
+            !llvm::compression::zstd::isAvailable() ||
+            static_cast<int>(Bin->getBufferSize()) < OffloadCompressThreshold) {
+          Fbin = addDeviceImageToModule(
+              ArrayRef<char>(Bin->getBufferStart(), Bin->getBufferSize()),
+              Twine(OffloadKindTag) + Twine(ImgId) + Twine(".data"), Kind,
+              Img.Tgt);
+        } else {
+
+          // Compress the image using zstd.
+          SmallVector<uint8_t, 512> CompressedBuffer;
+          llvm::compression::zstd::compress(
+              ArrayRef<unsigned char>(
+                  (const unsigned char *)(Bin->getBufferStart()),
+                  Bin->getBufferSize()),
+              CompressedBuffer, OffloadCompressLevel);
+
+          if (Verbose)
+            errs() << "[Compression] Original image size: "
+                   << Bin->getBufferSize() << "\n"
+                   << "[Compression] Compressed image size: "
+                   << CompressedBuffer.size() << "\n"
+                   << "[Compression] Compression level used: "
+                   << OffloadCompressLevel << "\n";
+
+          // Add the compressed image to the module.
+          Fbin = addDeviceImageToModule(
+              ArrayRef<char>((const char *)CompressedBuffer.data(),
+                             CompressedBuffer.size()),
+              Twine(OffloadKindTag) + Twine(ImgId) + Twine(".data"), Kind,
+              Img.Tgt);
+
+          // Change image format to compressed_non.
+          Ffmt = ConstantInt::get(Type::getInt8Ty(C),
+                                  BinaryImageFormat::compressed_none);
+        }
       }
 
       if (Kind == OffloadKind::SYCL) {

llvm/cmake/modules/Findzstd.cmake

@@ -16,11 +16,12 @@ else()
   set(zstd_STATIC_LIBRARY_SUFFIX "\\${CMAKE_STATIC_LIBRARY_SUFFIX}$")
 endif()
 
-find_path(zstd_INCLUDE_DIR NAMES zstd.h)
-find_library(zstd_LIBRARY NAMES zstd zstd_static)
+find_path(zstd_INCLUDE_DIR NAMES zstd.h HINTS $ENV{ZSTD_ROOT}/include)
+find_library(zstd_LIBRARY NAMES zstd zstd_static HINTS $ENV{ZSTD_ROOT}/lib)
 find_library(zstd_STATIC_LIBRARY NAMES
   zstd_static
-  "${CMAKE_STATIC_LIBRARY_PREFIX}zstd${CMAKE_STATIC_LIBRARY_SUFFIX}")
+  "${CMAKE_STATIC_LIBRARY_PREFIX}zstd${CMAKE_STATIC_LIBRARY_SUFFIX}"
+  HINTS $ENV{ZSTD_ROOT}/lib)
 
 include(FindPackageHandleStandardArgs)
 find_package_handle_standard_args(

sycl/doc/UsersManual.md

@@ -195,6 +195,19 @@ and not recommended to use in production environment.
     which may or may not perform additional inlining.
     Default value is 225.
 
+**`--offload-compress`**
+
+    Enables device image compression for SYCL offloading. Device images
+    are compressed using `zstd` compression algorithm and only if their size
+    exceeds 512 bytes.
+    Default value is false.
+
+**`--offload-compression-level=<int>`**
+
+    `zstd` compression level used to compress device images when `--offload-
+    compress` is enabled.
+    The default value is 10.
+
 ## Target toolchain options
 
 **`-Xsycl-target-backend=<T> "options"`**

sycl/source/CMakeLists.txt

@@ -69,6 +69,15 @@ function(add_sycl_rt_library LIB_NAME LIB_OBJ_NAME)
     target_link_libraries(${LIB_NAME} PRIVATE ${ARG_XPTI_LIB})
   endif()
 
+  # TODO: Remove debug print. Need this to figure out why zstd is not
+  # being found on Windows CI machines.
+  set(CMAKE_FIND_DEBUG_MODE 1)
+
+  # Need zstd for device image compression.
+  find_package(zstd REQUIRED)
+  target_link_libraries(${LIB_NAME} PRIVATE ${zstd_STATIC_LIBRARY})
+  target_include_directories(${LIB_OBJ_NAME} PRIVATE ${zstd_INCLUDE_DIR})
+
   target_include_directories(${LIB_OBJ_NAME} PRIVATE ${BOOST_UNORDERED_INCLUDE_DIRS})
 
   # ur_win_proxy_loader

sycl/source/detail/compiler.hpp

@@ -115,7 +115,8 @@ enum sycl_device_binary_type : uint8_t {
   SYCL_DEVICE_BINARY_TYPE_NONE = 0,   // undetermined
   SYCL_DEVICE_BINARY_TYPE_NATIVE = 1, // specific to a device
   SYCL_DEVICE_BINARY_TYPE_SPIRV = 2,
-  SYCL_DEVICE_BINARY_TYPE_LLVMIR_BITCODE = 3
+  SYCL_DEVICE_BINARY_TYPE_LLVMIR_BITCODE = 3,
+  SYCL_DEVICE_BINARY_TYPE_COMPRESSED_NONE = 4
 };
 
 // Device binary descriptor version supported by this library.

sycl/source/detail/compression.hpp

@@ -0,0 +1,136 @@
+//==---------- compression.hpp --------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#pragma once
+
+#include <iostream>
+#include <memory>
+#include <zstd.h>
+
+#define ZSTD_CONTENTSIZE_UNKNOWN (0ULL - 1)
+#define ZSTD_CONTENTSIZE_ERROR (0ULL - 2)
+
+namespace sycl {
+inline namespace _V1 {
+namespace detail {
+
+// Singleton class to handle ZSTD compression and decompression.
+class ZSTDCompressor {
+private:
+  // Initialize ZSTD context and error code.
+  ZSTDCompressor() {
+    m_ZSTD_compression_ctx = static_cast<void *>(ZSTD_createCCtx());
+    m_ZSTD_decompression_ctx = static_cast<void *>(ZSTD_createDCtx());
+
+    if (!m_ZSTD_compression_ctx || !m_ZSTD_decompression_ctx) {
+      std::cerr << "Error creating ZSTD contexts. \n";
+    }
+
+    m_lastError = 0;
+  }
+
+  // Free ZSTD contexts.
+  ~ZSTDCompressor() {
+    ZSTD_freeCCtx(static_cast<ZSTD_CCtx *>(m_ZSTD_compression_ctx));
+    ZSTD_freeDCtx(static_cast<ZSTD_DCtx *>(m_ZSTD_decompression_ctx));
+  }
+
+  ZSTDCompressor(const ZSTDCompressor &) = delete;
+  ZSTDCompressor &operator=(const ZSTDCompressor &) = delete;
+
+  // Get the singleton instance of the ZSTDCompressor class.
+  static ZSTDCompressor &GetSingletonInstance() {
+    static ZSTDCompressor instance;
+    return instance;
+  }
+
+  // Public APIs
+public:
+  // Return 0 is last (de)compression was successful, otherwise return error
+  // code.
+  static int GetLastError() { return GetSingletonInstance().m_lastError; }
+
+  // Returns a string representation of the error code.
+  // If the error code is 0, it returns "No error detected".
+  static std::string GetErrorString(int code) {
+    return ZSTD_getErrorName(code);
+  }
+
+  // Blob (de)compression do not assume format/structure of the input buffer.
+  static std::unique_ptr<char> CompressBlob(const char *src, size_t srcSize,
+                                            size_t &dstSize, int level) {
+    auto &instance = GetSingletonInstance();
+
+    // Get maximum size of the compressed buffer and allocate it.
+    auto dstBufferSize = ZSTD_compressBound(srcSize);
+    auto dstBuffer = std::unique_ptr<char>(new char[dstBufferSize]);
+
+    // Compress the input buffer.
+    dstSize = ZSTD_compressCCtx(
+        static_cast<ZSTD_CCtx *>(instance.m_ZSTD_compression_ctx),
+        static_cast<void *>(dstBuffer.get()), dstBufferSize,
+        static_cast<const void *>(src), srcSize, level);
+
+    // Store the error code if compression failed.
+    if (ZSTD_isError(dstSize))
+      instance.m_lastError = dstSize;
+    else
+      instance.m_lastError = 0;
+
+    // Pass ownership of the buffer to the caller.
+    return dstBuffer;
+  }
+
+  static std::unique_ptr<unsigned char>
+  DecompressBlob(const char *src, size_t srcSize, size_t &dstSize) {
+    auto &instance = GetSingletonInstance();
+
+    // Size of decompressed image can be larger than what we can allocate
+    // on heap. In that case, we need to use streaming decompression.
+    // TODO: Throw if the decompression size is too large.
+    auto dstBufferSize = ZSTD_getFrameContentSize(src, srcSize);
+
+    if (dstBufferSize == ZSTD_CONTENTSIZE_UNKNOWN ||
+        dstBufferSize == ZSTD_CONTENTSIZE_ERROR) {
+
+      std::cerr << "Error determining size of uncompressed data\n";
+      dstSize = 0;
+      instance.m_lastError = dstBufferSize;
+      return nullptr;
+    }
+
+    // Allocate buffer for decompressed data.
+    auto dstBuffer =
+        std::unique_ptr<unsigned char>(new unsigned char[dstBufferSize]);
+
+    dstSize = ZSTD_decompressDCtx(
+        static_cast<ZSTD_DCtx *>(instance.m_ZSTD_decompression_ctx),
+        static_cast<void *>(dstBuffer.get()), dstBufferSize,
+        static_cast<const void *>(src), srcSize);
+
+    // In case of decompression error, return the error message and set dstSize
+    // to 0.
+    if (ZSTD_isError(dstSize)) {
+      instance.m_lastError = dstSize;
+      dstSize = 0;
+    }
+
+    // Pass ownership of the buffer to the caller.
+    return dstBuffer;
+  }
+
+  // Data fields
+private:
+  int m_lastError;
+  // ZSTD context. Reusing ZSTD context speeds up subsequent (de)compression.
+  // Storing as void* to avoid including ZSTD headers in this file.
+  void *m_ZSTD_compression_ctx;
+  void *m_ZSTD_decompression_ctx;
+};
+} // namespace detail
+} // namespace _V1
+} // namespace sycl