Skip to content

Commit fec9509

Browse files
yxsamliuyxsamliu
committed
[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.
1 parent 9e7aab9 commit fec9509

File tree

4 files changed

+64
-31
lines changed

4 files changed

+64
-31
lines changed

‎clang/docs/ClangOffloadBundler.rst

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,7 +518,10 @@ The compressed offload bundle begins with a header followed by the compressed bi
518518
This is a unique identifier to distinguish compressed offload bundles. The value is the string 'CCOB' (Compressed Clang Offload Bundle).
519519

520520
- **Version Number (16-bit unsigned int)**:
521-
This denotes the version of the compressed offload bundle format. The current version is `1`.
521+
This denotes the version of the compressed offload bundle format. The current version is `2`.
522+
523+
- **Total File Size (32-bit unsigned int)**:
524+
This is the total size (in bytes) of the file, including the header. Available in version 2 and above.
522525

523526
- **Compression Method (16-bit unsigned int)**:
524527
This field indicates the compression method used. The value corresponds to either `zlib` or `zstd`, represented as a 16-bit unsigned integer cast from the LLVM compression enumeration.

‎clang/include/clang/Driver/OffloadBundler.h

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ struct OffloadTargetInfo {
9898
// The format is as follows:
9999
// - Magic Number (4 bytes) - A constant "CCOB".
100100
// - Version (2 bytes)
101+
// - Total file size (4 bytes). Available in version 2 and above.
101102
// - Compression Method (2 bytes) - Uses the values from
102103
// llvm::compression::Format.
103104
// - Uncompressed Size (4 bytes).
@@ -108,14 +109,18 @@ class CompressedOffloadBundle {
108109
private:
109110
static inline const size_t MagicSize = 4;
110111
static inline const size_t VersionFieldSize = sizeof(uint16_t);
112+
static inline const size_t FileSizeFieldSize = sizeof(uint32_t);
111113
static inline const size_t MethodFieldSize = sizeof(uint16_t);
112-
static inline const size_t SizeFieldSize = sizeof(uint32_t);
113-
static inline const size_t HashFieldSize = 8;
114-
static inline const size_t HeaderSize = MagicSize + VersionFieldSize +
115-
MethodFieldSize + SizeFieldSize +
116-
HashFieldSize;
114+
static inline const size_t UncompressedSizeFieldSize = sizeof(uint32_t);
115+
static inline const size_t HashFieldSize = sizeof(uint64_t);
116+
static inline const size_t V1HeaderSize =
117+
MagicSize + VersionFieldSize + MethodFieldSize +
118+
UncompressedSizeFieldSize + HashFieldSize;
119+
static inline const size_t V2HeaderSize =
120+
MagicSize + VersionFieldSize + FileSizeFieldSize + MethodFieldSize +
121+
UncompressedSizeFieldSize + HashFieldSize;
117122
static inline const llvm::StringRef MagicNumber = "CCOB";
118-
static inline const uint16_t Version = 1;
123+
static inline const uint16_t Version = 2;
119124

120125
public:
121126
static llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>

‎clang/lib/Driver/OffloadBundler.cpp

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1010,11 +1010,17 @@ CompressedOffloadBundle::compress(llvm::compression::Params P,
10101010

10111011
uint16_t CompressionMethod = static_cast<uint16_t>(P.format);
10121012
uint32_t UncompressedSize = Input.getBuffer().size();
1013+
uint32_t TotalFileSize = MagicNumber.size() + sizeof(TotalFileSize) +
1014+
sizeof(Version) + sizeof(CompressionMethod) +
1015+
sizeof(UncompressedSize) + sizeof(TruncatedHash) +
1016+
CompressedBuffer.size();
10131017

10141018
SmallVector<char, 0> FinalBuffer;
10151019
llvm::raw_svector_ostream OS(FinalBuffer);
10161020
OS << MagicNumber;
10171021
OS.write(reinterpret_cast<const char *>(&Version), sizeof(Version));
1022+
OS.write(reinterpret_cast<const char *>(&TotalFileSize),
1023+
sizeof(TotalFileSize));
10181024
OS.write(reinterpret_cast<const char *>(&CompressionMethod),
10191025
sizeof(CompressionMethod));
10201026
OS.write(reinterpret_cast<const char *>(&UncompressedSize),
@@ -1034,6 +1040,8 @@ CompressedOffloadBundle::compress(llvm::compression::Params P,
10341040
(UncompressedSize / (1024.0 * 1024.0)) / CompressionTimeSeconds;
10351041

10361042
llvm::errs() << "Compressed bundle format version: " << Version << "\n"
1043+
<< "Total file size (including headers): "
1044+
<< formatWithCommas(TotalFileSize) << " bytes\n"
10371045
<< "Compression method used: " << MethodUsed << "\n"
10381046
<< "Compression level: " << P.level << "\n"
10391047
<< "Binary size before compression: "
@@ -1059,31 +1067,42 @@ CompressedOffloadBundle::decompress(const llvm::MemoryBuffer &Input,
10591067

10601068
StringRef Blob = Input.getBuffer();
10611069

1062-
if (Blob.size() < HeaderSize) {
1070+
if (Blob.size() < V1HeaderSize)
10631071
return llvm::MemoryBuffer::getMemBufferCopy(Blob);
1064-
}
1072+
10651073
if (llvm::identify_magic(Blob) !=
10661074
llvm::file_magic::offload_bundle_compressed) {
10671075
if (Verbose)
10681076
llvm::errs() << "Uncompressed bundle.\n";
10691077
return llvm::MemoryBuffer::getMemBufferCopy(Blob);
10701078
}
10711079

1080+
size_t CurrentOffset = MagicSize;
1081+
10721082
uint16_t ThisVersion;
1083+
memcpy(&ThisVersion, Blob.data() + CurrentOffset, sizeof(uint16_t));
1084+
CurrentOffset += VersionFieldSize;
1085+
1086+
uint32_t TotalFileSize;
1087+
if (ThisVersion >= 2) {
1088+
if (Blob.size() < V2HeaderSize)
1089+
return createStringError(inconvertibleErrorCode(),
1090+
"Compressed bundle header size too small");
1091+
memcpy(&TotalFileSize, Blob.data() + CurrentOffset, sizeof(uint32_t));
1092+
CurrentOffset += FileSizeFieldSize;
1093+
}
1094+
10731095
uint16_t CompressionMethod;
1096+
memcpy(&CompressionMethod, Blob.data() + CurrentOffset, sizeof(uint16_t));
1097+
CurrentOffset += MethodFieldSize;
1098+
10741099
uint32_t UncompressedSize;
1100+
memcpy(&UncompressedSize, Blob.data() + CurrentOffset, sizeof(uint32_t));
1101+
CurrentOffset += UncompressedSizeFieldSize;
1102+
10751103
uint64_t StoredHash;
1076-
memcpy(&ThisVersion, Input.getBuffer().data() + MagicNumber.size(),
1077-
sizeof(uint16_t));
1078-
memcpy(&CompressionMethod, Blob.data() + MagicSize + VersionFieldSize,
1079-
sizeof(uint16_t));
1080-
memcpy(&UncompressedSize,
1081-
Blob.data() + MagicSize + VersionFieldSize + MethodFieldSize,
1082-
sizeof(uint32_t));
1083-
memcpy(&StoredHash,
1084-
Blob.data() + MagicSize + VersionFieldSize + MethodFieldSize +
1085-
SizeFieldSize,
1086-
sizeof(uint64_t));
1104+
memcpy(&StoredHash, Blob.data() + CurrentOffset, sizeof(uint64_t));
1105+
CurrentOffset += HashFieldSize;
10871106

10881107
llvm::compression::Format CompressionFormat;
10891108
if (CompressionMethod ==
@@ -1102,7 +1121,7 @@ CompressedOffloadBundle::decompress(const llvm::MemoryBuffer &Input,
11021121
DecompressTimer.startTimer();
11031122

11041123
SmallVector<uint8_t, 0> DecompressedData;
1105-
StringRef CompressedData = Blob.substr(HeaderSize);
1124+
StringRef CompressedData = Blob.substr(CurrentOffset);
11061125
if (llvm::Error DecompressionError = llvm::compression::decompress(
11071126
CompressionFormat, llvm::arrayRefFromStringRef(CompressedData),
11081127
DecompressedData, UncompressedSize))
@@ -1135,8 +1154,11 @@ CompressedOffloadBundle::decompress(const llvm::MemoryBuffer &Input,
11351154
double DecompressionSpeedMBs =
11361155
(UncompressedSize / (1024.0 * 1024.0)) / DecompressionTimeSeconds;
11371156

1138-
llvm::errs() << "Compressed bundle format version: " << ThisVersion << "\n"
1139-
<< "Decompression method: "
1157+
llvm::errs() << "Compressed bundle format version: " << ThisVersion << "\n";
1158+
if (ThisVersion >= 2)
1159+
llvm::errs() << "Total file size (from header): "
1160+
<< formatWithCommas(TotalFileSize) << " bytes\n";
1161+
llvm::errs() << "Decompression method: "
11401162
<< (CompressionFormat == llvm::compression::Format::Zlib
11411163
? "zlib"
11421164
: "zstd")

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

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,19 +22,22 @@
2222
// Check compression/decompression of offload bundle.
2323
//
2424
// RUN: clang-offload-bundler -type=bc -targets=hip-amdgcn-amd-amdhsa--gfx900,hip-amdgcn-amd-amdhsa--gfx906 \
25-
// RUN: -input=%t.tgt1 -input=%t.tgt2 -output=%t.hip.bundle.bc -compress -verbose 2>&1 | \
26-
// RUN: FileCheck -check-prefix=COMPRESS %s
25+
// RUN: -input=%t.tgt1 -input=%t.tgt2 -output=%t.hip.bundle.bc -compress -verbose >%t.1.txt 2>&1
2726
// RUN: clang-offload-bundler -type=bc -list -input=%t.hip.bundle.bc | FileCheck -check-prefix=NOHOST %s
2827
// RUN: clang-offload-bundler -type=bc -targets=hip-amdgcn-amd-amdhsa--gfx900,hip-amdgcn-amd-amdhsa--gfx906 \
29-
// RUN: -output=%t.res.tgt1 -output=%t.res.tgt2 -input=%t.hip.bundle.bc -unbundle -verbose 2>&1 | \
30-
// RUN: FileCheck -check-prefix=DECOMPRESS %s
28+
// RUN: -output=%t.res.tgt1 -output=%t.res.tgt2 -input=%t.hip.bundle.bc -unbundle -verbose >%t.2.txt 2>&1
29+
// RUN: cat %t.1.txt %t.2.txt | FileCheck %s
3130
// RUN: diff %t.tgt1 %t.res.tgt1
3231
// RUN: diff %t.tgt2 %t.res.tgt2
3332
//
34-
// COMPRESS: Compression method used: zstd
35-
// COMPRESS: Compression level: 3
36-
// DECOMPRESS: Decompression method: zstd
37-
// DECOMPRESS: Hashes match: Yes
33+
// CHECK: Compressed bundle format version: 2
34+
// CHECK: Total file size (including headers): [[SIZE:[0-9]*]] bytes
35+
// CHECK: Compression method used: zstd
36+
// CHECK: Compression level: 3
37+
// CHECK: Compressed bundle format version: 2
38+
// CHECK: Total file size (from header): [[SIZE]] bytes
39+
// CHECK: Decompression method: zstd
40+
// CHECK: Hashes match: Yes
3841
// NOHOST-NOT: host-
3942
// NOHOST-DAG: hip-amdgcn-amd-amdhsa--gfx900
4043
// NOHOST-DAG: hip-amdgcn-amd-amdhsa--gfx906

0 commit comments

Comments
 (0)