[SYCL] Support --offload-compress in clang-linker-wrapper

clang/test/Driver/sycl-offload-new-driver-compression.cpp

@@ -0,0 +1,26 @@
+/// Check that '--offload-compress' and '--offload-compression-level=' are
+/// forwarded to clang-linker-wrapper under the new offloading driver via the
+/// existing addOffloadCompressArgs() plumbing (as '--compress' and
+/// '--compression-level='). The wrapper consumes them in
+/// wrapSYCLBinariesFromFile to zstd-compress each SYCL device image and tag
+/// it with BIF_Compressed.
+
+// RUN: %clangxx -### -fsycl --offload-new-driver --offload-compress \
+// RUN:     --offload-compression-level=3 %s 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-COMPRESS
+// CHECK-COMPRESS: {{.*}}clang-linker-wrapper{{.*}}"--compress"{{.*}}"--compression-level=3"
+
+// RUN: %clangxx -### -fsycl --offload-new-driver %s 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-NO-COMPRESS
+// CHECK-NO-COMPRESS-NOT: {{.*}}clang-linker-wrapper{{.*}}"--compress"
+
+// --no-offload-compress overrides an earlier --offload-compress.
+// RUN: %clangxx -### -fsycl --offload-new-driver --offload-compress \
+// RUN:     --no-offload-compress %s 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-NO-COMPRESS
+
+// ...and vice versa: a later --offload-compress wins over --no-offload-compress.
+// RUN: %clangxx -### -fsycl --offload-new-driver --no-offload-compress \
+// RUN:     --offload-compress %s 2>&1 \
+// RUN:   | FileCheck %s --check-prefix=CHECK-COMPRESS-LAST
+// CHECK-COMPRESS-LAST: {{.*}}clang-linker-wrapper{{.*}}"--compress"

clang/test/OffloadTools/clang-linker-wrapper/linker-wrapper-zstd.cpp

@@ -0,0 +1,72 @@
+// REQUIRES: zstd && system-linux && x86-registered-target
+
+// clang-linker-wrapper compression test: checks that the wrapper compresses
+// SYCL device images when --compress is set, tags them with
+// SYCLBinaryImageFormat::BIF_Compressed (int8 value 4) in the emitted wrapper
+// module, honors --compression-level=, and reports invalid values.
+
+// Prepare test data. The bitcode is over the 512-byte compression threshold,
+// so compression actually runs.
+// RUN: %clang_cc1 -fsycl-is-device -disable-llvm-passes -triple=spir64-unknown-unknown %s -emit-llvm-bc -o %t.device.bc
+// RUN: llvm-offload-binary -o %t.fat --image=file=%t.device.bc,kind=sycl,triple=spir64-unknown-unknown
+// RUN: %clang_cc1 %s -triple=x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.fat
+// RUN: touch %t.devicelib.bc
+
+// With --compress and --wrapper-verbose the compression path fires and the
+// wrapped image's Format slot in %__sycl.tgt_device_image is i8 4
+// (BIF_Compressed).
+// RUN: clang-linker-wrapper --print-wrapped-module --host-triple=x86_64-unknown-linux-gnu \
+// RUN:   --bitcode-library=spir64-unknown-unknown=%t.devicelib.bc \
+// RUN:   --compress --compression-level=9 --wrapper-verbose \
+// RUN:   %t.o -o %t.out --linker-path=/usr/bin/ld 2>&1 \
+// RUN: | FileCheck %s --check-prefix=CHECK-COMPRESS
+
+// Capture the original and compressed sizes.
+// CHECK-COMPRESS: [Compression] Original image size: [[#ORIG:]]
+// CHECK-COMPRESS: [Compression] Compressed image size: [[#COMP:]]
+// CHECK-COMPRESS: [Compression] Compression level used: 9
+//
+// Assert ORIG - COMP > 0. FileCheck's numeric matching form only supports
+// `==`, but sub() rejects underflow at expression-evaluation time. We attach
+// sub(ORIG, COMP) to a CHECK-NOT pattern that begins with a sentinel string
+// never present in the output:
+//   * when ORIG > COMP, sub() succeeds; the substituted pattern is
+//     "COMPRESSION_SIZE_CHECK<some-number>", which doesn't appear, so the
+//     CHECK-NOT is satisfied.
+//   * when COMP >= ORIG, sub() underflows and FileCheck fails the pattern.
+// CHECK-COMPRESS-NOT: COMPRESSION_SIZE_CHECK[[#sub(ORIG, COMP)]]
+// CHECK-COMPRESS: @.sycl_offloading.device_images = internal unnamed_addr constant [1 x %__sycl.tgt_device_image] [%__sycl.tgt_device_image { i16 {{[0-9]+}}, i8 4, i8 4,
+
+// Without --compress the image is left untagged (Format = BIF_None = 0) and
+// no [Compression] verbose lines are emitted.
+// RUN: clang-linker-wrapper --print-wrapped-module --host-triple=x86_64-unknown-linux-gnu \
+// RUN:   --bitcode-library=spir64-unknown-unknown=%t.devicelib.bc \
+// RUN:   --wrapper-verbose \
+// RUN:   %t.o -o %t.out --linker-path=/usr/bin/ld 2>&1 \
+// RUN: | FileCheck %s --check-prefix=CHECK-NO-COMPRESS
+
+// CHECK-NO-COMPRESS-NOT: [Compression]
+// CHECK-NO-COMPRESS: @.sycl_offloading.device_images = internal unnamed_addr constant [1 x %__sycl.tgt_device_image] [%__sycl.tgt_device_image { i16 {{[0-9]+}}, i8 4, i8 0,
+
+// A non-integer --compression-level= is diagnosed.
+// RUN: not clang-linker-wrapper --host-triple=x86_64-unknown-linux-gnu \
+// RUN:   --bitcode-library=spir64-unknown-unknown=%t.devicelib.bc \
+// RUN:   --compress --compression-level=notanumber \
+// RUN:   %t.o -o %t.out --linker-path=/usr/bin/ld 2>&1 \
+// RUN: | FileCheck %s --check-prefix=CHECK-BAD-LEVEL
+
+// CHECK-BAD-LEVEL: invalid value for --offload-compression-level=: 'notanumber'
+
+template <typename t, typename Func>
+__attribute__((sycl_kernel)) void kernel(const Func &func) {
+    func();
+}
+
+extern "C" {
+void __sycl_register_lib(void *) {}
+void __sycl_unregister_lib(void *) {}
+}
+
+int main() {
+    kernel<class fake_kernel>([](){});
+}

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

@@ -1107,6 +1107,41 @@ static Expected<StringRef> runAOTCompile(StringRef InputFile,
                            "Unsupported SYCL Triple and Arch");
 }
 
+/// Compress each SYCL device image in-place when --compress is set,
+/// tagging its Format as BIF_Compressed. --compress and --compression-level=
+/// are the same flags HIP forwards to clang-offload-bundler;
+static Error compressImages(SmallVectorImpl<offloading::SYCLImage> &Images,
+                            const ArgList &Args) {
+  if (!Args.hasArg(OPT_compress))
+    return Error::success();
+
+  int Level = llvm::offloading::DefaultSYCLCompressionLevel;
+  if (auto *A = Args.getLastArg(OPT_compression_level_eq))
+    if (StringRef(A->getValue()).getAsInteger(10, Level))
+      return createStringError(
+          "invalid value for --offload-compression-level=: '%s'",
+          A->getValue());
+
+  for (auto &Image : Images) {
+    SmallVector<uint8_t, 0> CompressedBytes;
+    Expected<bool> DidCompressOrErr = offloading::compressSYCLDeviceImage(
+        ArrayRef<uint8_t>(
+            reinterpret_cast<const uint8_t *>(Image.Image->getBufferStart()),
+            Image.Image->getBufferSize()),
+        CompressedBytes, Level, /*Threshold=*/512, Verbose);
+    if (!DidCompressOrErr)
+      return DidCompressOrErr.takeError();
+    if (!*DidCompressOrErr)
+      continue;
+    Image.Image = MemoryBuffer::getMemBufferCopy(
+        StringRef(reinterpret_cast<const char *>(CompressedBytes.data()),
+                  CompressedBytes.size()),
+        Image.Image->getBufferIdentifier());
+    Image.Format = offloading::SYCLBinaryImageFormat::BIF_Compressed;
+  }
+  return Error::success();
+}
+
 /// Reads device images from the given \p InputFile and wraps them
 /// in one LLVM IR Module as a constant data.
 ///
@@ -1178,6 +1213,9 @@ wrapSYCLBinariesFromFile(ArrayRef<module_split::SplitModule> SplitModules,
                         ImageTarget, SI.CompileOptions, SI.LinkOptions);
   }
 
+  if (Error E = compressImages(Images, Args))
+    return std::move(E);
+
   LLVMContext C;
   Module M("offload.wrapper.object", C);
   M.setTargetTriple(Triple(

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

@@ -1,4 +1,4 @@
-set(LLVM_LINK_COMPONENTS BitWriter BitReader Core Object Support TransformUtils TargetParser)
+set(LLVM_LINK_COMPONENTS BitWriter BitReader Core FrontendOffloading Object Support TransformUtils TargetParser)
 
 add_clang_tool(clang-offload-wrapper
   ClangOffloadWrapper.cpp