[ClangOffloadBundler] Add file size to header

__hipRegisterFatBinary only accepts one pointer argument. It is
expected to get the fat binary size from the header.

This patch adds a file size field to the header of the compressed
bundle.

Author: yxsamliu

clang/docs/ClangOffloadBundler.rst

@@ -517,6 +517,9 @@ The compressed offload bundle begins with a header followed by the compressed bi
 - **Magic Number (4 bytes)**:
     This is a unique identifier to distinguish compressed offload bundles. The value is the string 'CCOB' (Compressed Clang Offload Bundle).
 
+- **Total File Size (32-bit unsigned int)**:
+    This is the total size (in bytes) of the file, including the header.
+
 - **Version Number (16-bit unsigned int)**:
     This denotes the version of the compressed offload bundle format. The current version is `1`.
 

clang/include/clang/Driver/OffloadBundler.h

@@ -97,6 +97,7 @@ struct OffloadTargetInfo {
 //
 // The format is as follows:
 // - Magic Number (4 bytes) - A constant "CCOB".
+// - Total file size (4 bytes).
 // - Version (2 bytes)
 // - Compression Method (2 bytes) - Uses the values from
 // llvm::compression::Format.
@@ -107,13 +108,14 @@ struct OffloadTargetInfo {
 class CompressedOffloadBundle {
 private:
   static inline const size_t MagicSize = 4;
+  static inline const size_t FileSizeFieldSize = sizeof(uint32_t);
   static inline const size_t VersionFieldSize = sizeof(uint16_t);
   static inline const size_t MethodFieldSize = sizeof(uint16_t);
-  static inline const size_t SizeFieldSize = sizeof(uint32_t);
-  static inline const size_t HashFieldSize = 8;
-  static inline const size_t HeaderSize = MagicSize + VersionFieldSize +
-                                          MethodFieldSize + SizeFieldSize +
-                                          HashFieldSize;
+  static inline const size_t UncompressedSizeFieldSize = sizeof(uint32_t);
+  static inline const size_t HashFieldSize = sizeof(uint64_t);
+  static inline const size_t HeaderSize =
+      MagicSize + FileSizeFieldSize + VersionFieldSize + MethodFieldSize +
+      UncompressedSizeFieldSize + HashFieldSize;
   static inline const llvm::StringRef MagicNumber = "CCOB";
   static inline const uint16_t Version = 1;
 

clang/lib/Driver/OffloadBundler.cpp

@@ -1010,10 +1010,16 @@ CompressedOffloadBundle::compress(llvm::compression::Params P,
 
   uint16_t CompressionMethod = static_cast<uint16_t>(P.format);
   uint32_t UncompressedSize = Input.getBuffer().size();
+  uint32_t TotalFileSize = MagicNumber.size() + sizeof(TotalFileSize) +
+                           sizeof(Version) + sizeof(CompressionMethod) +
+                           sizeof(UncompressedSize) + sizeof(TruncatedHash) +
+                           CompressedBuffer.size();
 
   SmallVector<char, 0> FinalBuffer;
   llvm::raw_svector_ostream OS(FinalBuffer);
   OS << MagicNumber;
+  OS.write(reinterpret_cast<const char *>(&TotalFileSize),
+           sizeof(TotalFileSize));
   OS.write(reinterpret_cast<const char *>(&Version), sizeof(Version));
   OS.write(reinterpret_cast<const char *>(&CompressionMethod),
            sizeof(CompressionMethod));
@@ -1033,7 +1039,9 @@ CompressedOffloadBundle::compress(llvm::compression::Params P,
     double CompressionSpeedMBs =
         (UncompressedSize / (1024.0 * 1024.0)) / CompressionTimeSeconds;
 
-    llvm::errs() << "Compressed bundle format version: " << Version << "\n"
+    llvm::errs() << "Total file size (including headers): "
+                 << formatWithCommas(TotalFileSize) << " bytes\n"
+                 << "Compressed bundle format version: " << Version << "\n"
                  << "Compression method used: " << MethodUsed << "\n"
                  << "Compression level: " << P.level << "\n"
                  << "Binary size before compression: "
@@ -1069,21 +1077,26 @@ CompressedOffloadBundle::decompress(const llvm::MemoryBuffer &Input,
     return llvm::MemoryBuffer::getMemBufferCopy(Blob);
   }
 
+  size_t CurrentOffset = MagicSize;
+  uint32_t TotalFileSize;
+  memcpy(&TotalFileSize, Blob.data() + CurrentOffset, sizeof(uint32_t));
+  CurrentOffset += FileSizeFieldSize;
+
   uint16_t ThisVersion;
+  memcpy(&ThisVersion, Blob.data() + CurrentOffset, sizeof(uint16_t));
+  CurrentOffset += VersionFieldSize;
+
   uint16_t CompressionMethod;
+  memcpy(&CompressionMethod, Blob.data() + CurrentOffset, sizeof(uint16_t));
+  CurrentOffset += MethodFieldSize;
+
   uint32_t UncompressedSize;
+  memcpy(&UncompressedSize, Blob.data() + CurrentOffset, sizeof(uint32_t));
+  CurrentOffset += UncompressedSizeFieldSize;
+
   uint64_t StoredHash;
-  memcpy(&ThisVersion, Input.getBuffer().data() + MagicNumber.size(),
-         sizeof(uint16_t));
-  memcpy(&CompressionMethod, Blob.data() + MagicSize + VersionFieldSize,
-         sizeof(uint16_t));
-  memcpy(&UncompressedSize,
-         Blob.data() + MagicSize + VersionFieldSize + MethodFieldSize,
-         sizeof(uint32_t));
-  memcpy(&StoredHash,
-         Blob.data() + MagicSize + VersionFieldSize + MethodFieldSize +
-             SizeFieldSize,
-         sizeof(uint64_t));
+  memcpy(&StoredHash, Blob.data() + CurrentOffset, sizeof(uint64_t));
+  CurrentOffset += HashFieldSize;
 
   llvm::compression::Format CompressionFormat;
   if (CompressionMethod ==
@@ -1135,7 +1148,9 @@ CompressedOffloadBundle::decompress(const llvm::MemoryBuffer &Input,
     double DecompressionSpeedMBs =
         (UncompressedSize / (1024.0 * 1024.0)) / DecompressionTimeSeconds;
 
-    llvm::errs() << "Compressed bundle format version: " << ThisVersion << "\n"
+    llvm::errs() << "Total file size (from header): "
+                 << formatWithCommas(TotalFileSize) << " bytes\n"
+                 << "Compressed bundle format version: " << ThisVersion << "\n"
                  << "Decompression method: "
                  << (CompressionFormat == llvm::compression::Format::Zlib
                          ? "zlib"

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

@@ -22,19 +22,20 @@
 // Check compression/decompression of offload bundle.
 //
 // RUN: clang-offload-bundler -type=bc -targets=hip-amdgcn-amd-amdhsa--gfx900,hip-amdgcn-amd-amdhsa--gfx906 \
-// RUN:   -input=%t.tgt1 -input=%t.tgt2 -output=%t.hip.bundle.bc -compress -verbose 2>&1 | \
-// RUN:   FileCheck -check-prefix=COMPRESS %s
+// RUN:   -input=%t.tgt1 -input=%t.tgt2 -output=%t.hip.bundle.bc -compress -verbose >%t.1.txt 2>&1
 // RUN: clang-offload-bundler -type=bc -list -input=%t.hip.bundle.bc | FileCheck -check-prefix=NOHOST %s
 // RUN: clang-offload-bundler -type=bc -targets=hip-amdgcn-amd-amdhsa--gfx900,hip-amdgcn-amd-amdhsa--gfx906 \
-// RUN:   -output=%t.res.tgt1 -output=%t.res.tgt2 -input=%t.hip.bundle.bc -unbundle -verbose 2>&1 | \
-// RUN:   FileCheck -check-prefix=DECOMPRESS %s
+// RUN:   -output=%t.res.tgt1 -output=%t.res.tgt2 -input=%t.hip.bundle.bc -unbundle -verbose >%t.2.txt 2>&1
+// RUN: cat %t.1.txt %t.2.txt | FileCheck %s
 // RUN: diff %t.tgt1 %t.res.tgt1
 // RUN: diff %t.tgt2 %t.res.tgt2
 //
-// COMPRESS: Compression method used: zstd
-// COMPRESS: Compression level: 3 
-// DECOMPRESS: Decompression method: zstd
-// DECOMPRESS: Hashes match: Yes
+// CHECK: Total file size (including headers): [[SIZE:[0-9]*]] bytes
+// CHECK: Compression method used: zstd
+// CHECK: Compression level: 3
+// CHECK: Total file size (from header): [[SIZE]] bytes
+// CHECK: Decompression method: zstd
+// CHECK: Hashes match: Yes
 // NOHOST-NOT: host-
 // NOHOST-DAG: hip-amdgcn-amd-amdhsa--gfx900
 // NOHOST-DAG: hip-amdgcn-amd-amdhsa--gfx906