[SYCL][Clang] Add support for device image compression (#15124)

This PR adds support for device image compression for the old offloading
model. I'll make another follow-up PR to extend support for the new
offload model.

**Design summary**:
```
ZSTD (compression algo)   ----> LLVMSupport (Interface)  ------> clang-offload-wrapper (For compression)
 |
 ----------------------------------------------------- --------> SYCL RT (For decompression)
```
This PR introduces ZSTD (https://github.com/facebook/zstd) as a 3rd
party dependency of DPCPP. **Similar to upstream LLVM, we expect user to
have `zstd-dev` package installed on their machine** - we won't be
installing zstd from sources.

**How to use**
To compress device images, add `--offload-compress` CLI option to your
clang invocation. Note that we compress device images only if the size
of device images exceeds a threshold, which is 512 bytes by default.
Moreover, by default, we use ZSTD level 10 for compression. ZSTD
compression levels provides a tradeoff between (de)compression time and
compression ratio, and the compression level can be changed using
`--offload-compression-level=<int>` CLI option.

Author: uditagarwal97

buildbot/configure.py

@@ -177,6 +177,8 @@ 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",
+        "-DLLVM_USE_STATIC_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

@@ -10174,6 +10174,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/test/Driver/sycl-offload-wrapper-compression.cpp

@@ -0,0 +1,14 @@
+///
+/// Check if '--offload-compress' and '--offload-compression-level' CLI
+/// options are passed to the clang-offload-wrapper.
+///
+
+// RUN: %clangxx -### -fsycl --offload-compress --offload-compression-level=3 %s 2>&1 | FileCheck %s --check-prefix=CHECK-COMPRESS
+// CHECK-COMPRESS: {{.*}}clang-offload-wrapper{{.*}}"-offload-compress"{{.*}}"-offload-compression-level=3"{{.*}}
+
+// Make sure that the compression options are not passed when --offload-compress is not set.
+// RUN: %clangxx -### -fsycl %s 2>&1 | FileCheck %s --check-prefix=CHECK-NO-COMPRESS
+// RUN: %clangxx -### -fsycl --offload-compression-level=3 %s 2>&1 | FileCheck %s --check-prefix=CHECK-NO-COMPRESS
+
+// CHECK-NO-COMPRESS-NOT: {{.*}}clang-offload-wrapper{{.*}}"-offload-compress"{{.*}}
+// CHECK-NO-COMPRESS-NOT: {{.*}}clang-offload-wrapper{{.*}}"-offload-compression-level=3"{{.*}}

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,66 @@ 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, throw an error.
+        if (OffloadCompressDevImgs && !llvm::compression::zstd::isAvailable()) {
+          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
+        // 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;
+#if LLVM_ENABLE_EXCEPTIONS
+          try {
+#endif
+            llvm::compression::zstd::compress(
+                ArrayRef<unsigned char>(
+                    (const unsigned char *)(Bin->getBufferStart()),
+                    Bin->getBufferSize()),
+                CompressedBuffer, OffloadCompressLevel);
+#if LLVM_ENABLE_EXCEPTIONS
+          } catch (const std::exception &ex) {
+            return createStringError(inconvertibleErrorCode(),
+                                     std::string("Failed to compress the device image: \n") +
+                                     std::string(ex.what()));
+          }
+#endif
+          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_none.
+          Ffmt = ConstantInt::get(Type::getInt8Ty(C),
+                                  BinaryImageFormat::compressed_none);
+        }
       }
 
       if (Kind == OffloadKind::SYCL) {

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,13 @@ function(add_sycl_rt_library LIB_NAME LIB_OBJ_NAME)
     target_link_libraries(${LIB_NAME} PRIVATE ${ARG_XPTI_LIB})
   endif()
 
+  if (NOT LLVM_ENABLE_ZSTD)
+    target_compile_definitions(${LIB_OBJ_NAME} PRIVATE SYCL_RT_ZSTD_NOT_AVAIABLE)
+  else()
+    target_link_libraries(${LIB_NAME} PRIVATE ${zstd_STATIC_LIBRARY})
+    target_include_directories(${LIB_OBJ_NAME} PRIVATE ${zstd_INCLUDE_DIR})
+  endif()
+
   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.