[llvm-branch-commits] [llvm] [Offloading] Add support for compressed OffloadBinary types (PR #222774)
Joseph Huber via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Thu Sep 10 14:08:26 PDT 2026
https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/222774
>From be78c5ad8a7a80d015ea2eb90d9df2a452df30d2 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Thu, 10 Sep 2026 15:58:45 -0500
Subject: [PATCH] [Offloading] Add support for compressed OffloadBinary types
Summary:
Offload binaries are used to store many heterogenous architectures into
a singel offloading blob. These lists can get very large so this PR adds
the option to compress them with the LLVM provided compression
libraries.
The implementation is quite simple, we simply compress all the buffers
after the header into a single compressed blob, then re-construct the
header. Extracting is the reverse.
The biggest change is that the offload binary now **owns** the memory,
whereas before we simply took a reference to it. This is necessary
because the decompression must create new memory compared to what the
user provided. This adds an extra copy internally, but it also
simplifies the V2 additions.
This does not wire up any clang/HIP support, just providing the
functionality.
---
llvm/include/llvm/Object/OffloadBinary.h | 18 +-
llvm/include/llvm/ObjectYAML/OffloadYAML.h | 5 +
llvm/lib/Object/OffloadBinary.cpp | 175 +++++++++++++-----
llvm/lib/ObjectYAML/OffloadEmitter.cpp | 15 +-
llvm/lib/ObjectYAML/OffloadYAML.cpp | 7 +
llvm/test/ObjectYAML/Offload/compressed.yaml | 30 +++
.../ObjectYAML/Offload/malformed-version.yaml | 2 +-
.../llvm-objdump/Offloading/compressed.test | 41 ++++
.../tools/llvm-offload-binary/compress.test | 19 ++
llvm/tools/llvm-objdump/OffloadDump.cpp | 13 +-
llvm/tools/llvm-objdump/OffloadDump.h | 2 +-
llvm/tools/llvm-objdump/llvm-objdump.cpp | 5 +-
.../llvm-offload-binary.cpp | 36 +++-
llvm/tools/obj2yaml/offload2yaml.cpp | 26 ++-
llvm/unittests/Object/OffloadingTest.cpp | 68 +++++++
15 files changed, 399 insertions(+), 63 deletions(-)
create mode 100644 llvm/test/ObjectYAML/Offload/compressed.yaml
create mode 100644 llvm/test/tools/llvm-objdump/Offloading/compressed.test
create mode 100644 llvm/test/tools/llvm-offload-binary/compress.test
diff --git a/llvm/include/llvm/Object/OffloadBinary.h b/llvm/include/llvm/Object/OffloadBinary.h
index ab664800567e4..9afe98815afb0 100644
--- a/llvm/include/llvm/Object/OffloadBinary.h
+++ b/llvm/include/llvm/Object/OffloadBinary.h
@@ -17,11 +17,13 @@
#ifndef LLVM_OBJECT_OFFLOADBINARY_H
#define LLVM_OBJECT_OFFLOADBINARY_H
+#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Object/Binary.h"
#include "llvm/Support/Compiler.h"
+#include "llvm/Support/Compression.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/MemoryBuffer.h"
#include <memory>
@@ -74,7 +76,7 @@ class OffloadBinary : public Binary {
using string_iterator_range = iterator_range<string_iterator>;
/// The current version of the binary used for backwards compatibility.
- static const uint32_t Version = 2;
+ static const uint32_t Version = 3;
/// The offloading metadata that will be serialized to a memory buffer.
struct OffloadingImage {
@@ -91,6 +93,7 @@ class OffloadBinary : public Binary {
uint64_t Size; // Size in bytes of this entire binary.
uint64_t EntriesOffset; // Offset in bytes to the start of entries block.
uint64_t EntriesCount; // Number of metadata entries in the binary.
+ uint64_t InflatedSize; // Original size of the binary if compressed.
};
struct Entry {
@@ -137,6 +140,10 @@ class OffloadBinary : public Binary {
LLVM_ABI static SmallString<0>
write(ArrayRef<OffloadingImage> OffloadingData);
+ /// Serialize \p OffloadingData to a compressed binary with \p Compress.
+ LLVM_ABI static Expected<SmallString<0>>
+ write(ArrayRef<OffloadingImage> OffloadingData, compression::Params Compress);
+
static uint64_t getAlignment() { return 8; }
ImageKind getImageKind() const { return TheEntry->TheImageKind; }
@@ -160,10 +167,11 @@ class OffloadBinary : public Binary {
static bool classof(const Binary *V) { return V->isOffloadFile(); }
private:
- OffloadBinary(MemoryBufferRef Source, const Header *TheHeader,
+ OffloadBinary(std::shared_ptr<MemoryBuffer> Owned, const Header *TheHeader,
const Entry *TheEntry, const uint64_t Index = 0)
- : Binary(Binary::ID_Offload, Source), Buffer(Source.getBufferStart()),
- TheHeader(TheHeader), TheEntry(TheEntry), Index(Index) {
+ : Binary(Binary::ID_Offload, *Owned), OwnedBuffer(std::move(Owned)),
+ Buffer(OwnedBuffer->getBufferStart()), TheHeader(TheHeader),
+ TheEntry(TheEntry), Index(Index) {
// StringEntryV1 and StringEntry have ABI compatible Key/ValueOffset fields,
// but different sizes, so we need to manually calculate offset.
const char *StringMapBegin = &Buffer[TheEntry->StringOffset];
@@ -187,6 +195,8 @@ class OffloadBinary : public Binary {
OffloadBinary(const OffloadBinary &Other) = delete;
+ /// Owned uncompressed binary. Shared between entries parsed from one blob.
+ std::shared_ptr<MemoryBuffer> OwnedBuffer;
/// Map from keys to offsets in the binary.
MapVector<StringRef, StringRef> StringData;
/// Raw pointer to the MemoryBufferRef for convenience.
diff --git a/llvm/include/llvm/ObjectYAML/OffloadYAML.h b/llvm/include/llvm/ObjectYAML/OffloadYAML.h
index 2ce335920f6ce..a4048c9dd80f2 100644
--- a/llvm/include/llvm/ObjectYAML/OffloadYAML.h
+++ b/llvm/include/llvm/ObjectYAML/OffloadYAML.h
@@ -41,6 +41,7 @@ struct Binary {
std::optional<uint64_t> Size;
std::optional<uint64_t> EntriesOffset;
std::optional<uint64_t> EntriesCount;
+ std::optional<compression::Format> Compression;
std::vector<Member> Members;
};
@@ -61,6 +62,10 @@ template <> struct ScalarEnumerationTraits<object::OffloadKind> {
LLVM_ABI static void enumeration(IO &IO, object::OffloadKind &Value);
};
+template <> struct ScalarEnumerationTraits<compression::Format> {
+ LLVM_ABI static void enumeration(IO &IO, compression::Format &Value);
+};
+
template <> struct MappingTraits<OffloadYAML::Binary> {
LLVM_ABI static void mapping(IO &IO, OffloadYAML::Binary &O);
};
diff --git a/llvm/lib/Object/OffloadBinary.cpp b/llvm/lib/Object/OffloadBinary.cpp
index 1774eafea1590..85d483dcc86eb 100644
--- a/llvm/lib/Object/OffloadBinary.cpp
+++ b/llvm/lib/Object/OffloadBinary.cpp
@@ -8,6 +8,8 @@
#include "llvm/Object/OffloadBinary.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/BinaryFormat/Magic.h"
#include "llvm/IR/Constants.h"
@@ -21,34 +23,21 @@
#include "llvm/Object/IRObjectFile.h"
#include "llvm/Object/ObjectFile.h"
#include "llvm/Support/Alignment.h"
+#include "llvm/Support/Compression.h"
+#include "llvm/Support/Error.h"
#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
#include "llvm/TargetParser/AMDGPUTargetParser.h"
+#include <cassert>
+#include <cstddef>
+#include <memory>
+
using namespace llvm;
using namespace llvm::object;
namespace {
-/// A MemoryBuffer that shares ownership of the underlying memory.
-/// This allows multiple OffloadBinary instances to share the same buffer.
-class SharedMemoryBuffer : public MemoryBuffer {
-public:
- SharedMemoryBuffer(std::shared_ptr<MemoryBuffer> Buf)
- : SharedBuf(std::move(Buf)) {
- init(SharedBuf->getBufferStart(), SharedBuf->getBufferEnd(),
- /*RequiresNullTerminator=*/false);
- }
-
- BufferKind getBufferKind() const override { return MemoryBuffer_Malloc; }
-
- StringRef getBufferIdentifier() const override {
- return SharedBuf->getBufferIdentifier();
- }
-
-private:
- const std::shared_ptr<MemoryBuffer> SharedBuf;
-};
-
/// Attempts to extract all the embedded device images contained inside the
/// buffer \p Contents. The buffer is expected to contain a valid offloading
/// binary format.
@@ -70,23 +59,18 @@ Error extractOffloadFiles(MemoryBufferRef Contents,
return HeaderOrErr.takeError();
const OffloadBinary::Header *Header = *HeaderOrErr;
- // Create a copy of original memory containing only the current binary.
- std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy(
- Buffer->getBuffer().take_front(Header->Size),
- Contents.getBufferIdentifier());
-
- auto BinariesOrErr = OffloadBinary::create(*BufferCopy);
+ MemoryBufferRef Slice(Buffer->getBuffer().take_front(Header->Size),
+ Contents.getBufferIdentifier());
+ auto BinariesOrErr = OffloadBinary::create(Slice);
if (!BinariesOrErr)
return BinariesOrErr.takeError();
- // Share ownership among multiple OffloadFiles.
- std::shared_ptr<MemoryBuffer> SharedBuffer =
- std::shared_ptr<MemoryBuffer>(std::move(BufferCopy));
-
for (auto &Binary : *BinariesOrErr) {
- std::unique_ptr<SharedMemoryBuffer> SharedBufferPtr =
- std::make_unique<SharedMemoryBuffer>(SharedBuffer);
- Binaries.emplace_back(std::move(Binary), std::move(SharedBufferPtr));
+ std::unique_ptr<MemoryBuffer> View = MemoryBuffer::getMemBuffer(
+ Binary->getMemoryBufferRef().getBuffer(),
+ Binary->getMemoryBufferRef().getBufferIdentifier(),
+ /*RequiresNullTerminator=*/false);
+ Binaries.emplace_back(std::move(Binary), std::move(View));
}
Offset += Header->Size;
@@ -197,11 +181,55 @@ Error extractFromArchive(const Archive &Library,
return Error::success();
}
+bool isCompressed(const OffloadBinary::Header &Header) {
+ return Header.Version >= 3 && Header.InflatedSize != 0;
+}
+
+Expected<std::unique_ptr<MemoryBuffer>>
+decompressOffloadBinary(MemoryBufferRef Buf) {
+ const auto *Header =
+ reinterpret_cast<const OffloadBinary::Header *>(Buf.getBufferStart());
+ if (Header->EntriesOffset > Header->Size ||
+ Header->InflatedSize < Header->EntriesOffset)
+ return errorCodeToError(object_error::unexpected_eof);
+
+ // Get the compressed binary blob after the header.
+ StringRef Compressed = Buf.getBuffer()
+ .take_front(Header->Size)
+ .drop_front(Header->EntriesOffset);
+ uint64_t BodySize = Header->InflatedSize - Header->EntriesOffset;
+
+ compression::Format Format = identify_magic(Compressed) == file_magic::zstd
+ ? compression::Format::Zstd
+ : compression::Format::Zlib;
+ if (const char *Reason = compression::getReasonIfUnsupported(Format))
+ return createStringError(Reason);
+
+ SmallVector<uint8_t, 0> Body;
+ if (Error Err = compression::decompress(
+ Format, arrayRefFromStringRef(Compressed), Body, BodySize))
+ return std::move(Err);
+
+ // Restore the old header data for the newly uncompressed blob.
+ OffloadBinary::Header Restored = *Header;
+ Restored.Size = Restored.InflatedSize;
+ Restored.InflatedSize = 0;
+ if (Restored.EntriesOffset + Body.size() != Restored.Size)
+ return errorCodeToError(object_error::parse_failed);
+
+ SmallString<0> Out;
+ Out.reserve(Restored.Size);
+ Out.append(StringRef(reinterpret_cast<const char *>(&Restored),
+ Restored.EntriesOffset));
+ Out.append(toStringRef(Body));
+ return MemoryBuffer::getMemBufferCopy(Out, Buf.getBufferIdentifier());
+}
+
} // namespace
Expected<const OffloadBinary::Header *>
OffloadBinary::extractHeader(MemoryBufferRef Buf) {
- if (Buf.getBufferSize() < sizeof(Header) + sizeof(Entry))
+ if (Buf.getBufferSize() < sizeof(Header))
return errorCodeToError(object_error::parse_failed);
// Check for 0x10FF1OAD magic bytes.
@@ -217,15 +245,21 @@ OffloadBinary::extractHeader(MemoryBufferRef Buf) {
if (TheHeader->Version == 0 || TheHeader->Version > OffloadBinary::Version)
return errorCodeToError(object_error::parse_failed);
- if (TheHeader->Size > Buf.getBufferSize() ||
- TheHeader->Size < sizeof(Entry) || TheHeader->Size < sizeof(Header))
+ if (TheHeader->Size > Buf.getBufferSize() || TheHeader->Size < sizeof(Header))
+ return errorCodeToError(object_error::unexpected_eof);
+
+ if (isCompressed(*TheHeader))
+ return TheHeader;
+
+ if (TheHeader->Size < sizeof(Entry))
return errorCodeToError(object_error::unexpected_eof);
uint64_t EntriesCount =
(TheHeader->Version == 1) ? 1 : TheHeader->EntriesCount;
uint64_t EntriesSize = sizeof(Entry) * EntriesCount;
if (TheHeader->EntriesOffset > TheHeader->Size - EntriesSize ||
- EntriesSize > TheHeader->Size - sizeof(Header))
+ // v1/v2 headers are 32 bytes; sizeof(Header) grew in v3.
+ EntriesSize > TheHeader->Size - offsetof(Header, InflatedSize))
return errorCodeToError(object_error::unexpected_eof);
return TheHeader;
@@ -233,20 +267,39 @@ OffloadBinary::extractHeader(MemoryBufferRef Buf) {
Expected<SmallVector<std::unique_ptr<OffloadBinary>>>
OffloadBinary::create(MemoryBufferRef Buf, std::optional<uint64_t> Index) {
- auto HeaderOrErr = OffloadBinary::extractHeader(Buf);
+ auto HeaderOrErr = extractHeader(Buf);
+ if (!HeaderOrErr)
+ return HeaderOrErr.takeError();
+ const Header *OnDisk = *HeaderOrErr;
+
+ // The binary data may be a compressed image.
+ std::shared_ptr<MemoryBuffer> Binary;
+ if (isCompressed(*OnDisk)) {
+ auto DecompressedOrErr = decompressOffloadBinary(Buf);
+ if (!DecompressedOrErr)
+ return DecompressedOrErr.takeError();
+ Binary = std::shared_ptr<MemoryBuffer>(std::move(*DecompressedOrErr));
+ } else {
+ Binary = std::shared_ptr<MemoryBuffer>(MemoryBuffer::getMemBufferCopy(
+ Buf.getBuffer().take_front(OnDisk->Size), Buf.getBufferIdentifier()));
+ }
+
+ // Owned is now an uncompressed OffloadBinary, parse it as before.
+ MemoryBufferRef Owned = *Binary;
+ HeaderOrErr = extractHeader(Owned);
if (!HeaderOrErr)
return HeaderOrErr.takeError();
const Header *TheHeader = *HeaderOrErr;
- const char *Start = Buf.getBufferStart();
+ const char *Start = Owned.getBufferStart();
const Entry *Entries =
reinterpret_cast<const Entry *>(&Start[TheHeader->EntriesOffset]);
auto validateEntry = [&](const Entry *TheEntry) -> Error {
- if (TheEntry->ImageOffset > Buf.getBufferSize() ||
- TheEntry->StringOffset > Buf.getBufferSize() ||
+ if (TheEntry->ImageOffset > Owned.getBufferSize() ||
+ TheEntry->StringOffset > Owned.getBufferSize() ||
TheEntry->StringOffset + TheEntry->NumStrings * sizeof(StringEntry) >
- Buf.getBufferSize())
+ Owned.getBufferSize())
return errorCodeToError(object_error::unexpected_eof);
return Error::success();
};
@@ -259,7 +312,8 @@ OffloadBinary::create(MemoryBufferRef Buf, std::optional<uint64_t> Index) {
if (auto Err = validateEntry(TheEntry))
return std::move(Err);
- Binaries.emplace_back(new OffloadBinary(Buf, TheHeader, TheEntry, *Index));
+ Binaries.emplace_back(
+ new OffloadBinary(Binary, TheHeader, TheEntry, *Index));
return std::move(Binaries);
}
@@ -269,7 +323,7 @@ OffloadBinary::create(MemoryBufferRef Buf, std::optional<uint64_t> Index) {
if (auto Err = validateEntry(TheEntry))
return std::move(Err);
- Binaries.emplace_back(new OffloadBinary(Buf, TheHeader, TheEntry, I));
+ Binaries.emplace_back(new OffloadBinary(Binary, TheHeader, TheEntry, I));
}
return std::move(Binaries);
@@ -305,7 +359,7 @@ SmallString<0> OffloadBinary::write(ArrayRef<OffloadingImage> OffloadingData) {
// Create the header and fill in the offsets. The entries will be directly
// placed after the header in memory. Align the size to the alignment of the
// header so this can be placed contiguously in a single section.
- Header TheHeader;
+ Header TheHeader{};
TheHeader.Size = alignTo(BinaryDataSize + TotalImagesSize, getAlignment());
TheHeader.EntriesOffset = sizeof(Header);
TheHeader.EntriesCount = EntriesCount;
@@ -363,6 +417,37 @@ SmallString<0> OffloadBinary::write(ArrayRef<OffloadingImage> OffloadingData) {
return Data;
}
+Expected<SmallString<0>>
+OffloadBinary::write(ArrayRef<OffloadingImage> OffloadingData,
+ compression::Params Compress) {
+ if (const char *Reason = compression::getReasonIfUnsupported(Compress.format))
+ return createStringError(Reason);
+
+ // Write the complete offloading binary as normal.
+ SmallString<0> Uncompressed = write(OffloadingData);
+ OffloadBinary::Header Header =
+ *reinterpret_cast<const OffloadBinary::Header *>(Uncompressed.data());
+
+ // Compress the entries after the header with the requested configuration.
+ StringRef Body = StringRef(Uncompressed).drop_front(Header.EntriesOffset);
+ SmallVector<uint8_t, 0> CompressedBuffer;
+ compression::compress(Compress, arrayRefFromStringRef(Body),
+ CompressedBuffer);
+
+ // Reset the header sizes and create the newly compressed binary.
+ Header.InflatedSize = Uncompressed.size();
+ Header.Size = Header.EntriesOffset + CompressedBuffer.size();
+
+ SmallString<0> Data;
+ Data.reserve(Header.Size);
+ raw_svector_ostream OS(Data);
+ OS << StringRef(reinterpret_cast<const char *>(&Header),
+ Header.EntriesOffset);
+ OS << toStringRef(CompressedBuffer);
+ assert(Header.Size == OS.tell() && "Size mismatch");
+ return Data;
+}
+
Error object::extractOffloadBinaries(MemoryBufferRef Buffer,
SmallVectorImpl<OffloadFile> &Binaries) {
file_magic Type = identify_magic(Buffer.getBuffer());
diff --git a/llvm/lib/ObjectYAML/OffloadEmitter.cpp b/llvm/lib/ObjectYAML/OffloadEmitter.cpp
index 51167bd812df3..7a8acc2114baf 100644
--- a/llvm/lib/ObjectYAML/OffloadEmitter.cpp
+++ b/llvm/lib/ObjectYAML/OffloadEmitter.cpp
@@ -9,6 +9,8 @@
#include "llvm/Object/OffloadBinary.h"
#include "llvm/ObjectYAML/OffloadYAML.h"
#include "llvm/ObjectYAML/yaml2obj.h"
+#include "llvm/Support/Compression.h"
+#include "llvm/Support/Error.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;
@@ -41,7 +43,18 @@ bool yaml2offload(Binary &Doc, raw_ostream &Out, ErrorHandler EH) {
}
// Copy the data to a new buffer so we can modify the bytes directly.
- auto Buffer = object::OffloadBinary::write(Images);
+ SmallString<0> Buffer;
+ if (Doc.Compression) {
+ Expected<SmallString<0>> CompressedOrErr = object::OffloadBinary::write(
+ Images, compression::Params(*Doc.Compression));
+ if (!CompressedOrErr) {
+ EH(toString(CompressedOrErr.takeError()));
+ return false;
+ }
+ Buffer = std::move(*CompressedOrErr);
+ } else {
+ Buffer = object::OffloadBinary::write(Images);
+ }
auto *TheHeader =
reinterpret_cast<object::OffloadBinary::Header *>(&Buffer[0]);
if (Doc.Version)
diff --git a/llvm/lib/ObjectYAML/OffloadYAML.cpp b/llvm/lib/ObjectYAML/OffloadYAML.cpp
index c0e0ed41aaca9..f61a2052073aa 100644
--- a/llvm/lib/ObjectYAML/OffloadYAML.cpp
+++ b/llvm/lib/ObjectYAML/OffloadYAML.cpp
@@ -44,6 +44,12 @@ void ScalarEnumerationTraits<object::OffloadKind>::enumeration(
IO.enumFallback<Hex16>(Value);
}
+void ScalarEnumerationTraits<compression::Format>::enumeration(
+ IO &IO, compression::Format &Value) {
+ IO.enumCase(Value, "zlib", compression::Format::Zlib);
+ IO.enumCase(Value, "zstd", compression::Format::Zstd);
+}
+
void MappingTraits<OffloadYAML::Binary>::mapping(IO &IO,
OffloadYAML::Binary &O) {
assert(!IO.getContext() && "The IO context is initialized already");
@@ -53,6 +59,7 @@ void MappingTraits<OffloadYAML::Binary>::mapping(IO &IO,
IO.mapOptional("Size", O.Size);
IO.mapOptional("EntriesOffset", O.EntriesOffset);
IO.mapOptional("EntriesCount", O.EntriesCount);
+ IO.mapOptional("Compression", O.Compression);
IO.mapRequired("Members", O.Members);
IO.setContext(nullptr);
}
diff --git a/llvm/test/ObjectYAML/Offload/compressed.yaml b/llvm/test/ObjectYAML/Offload/compressed.yaml
new file mode 100644
index 0000000000000..48cdc3669da78
--- /dev/null
+++ b/llvm/test/ObjectYAML/Offload/compressed.yaml
@@ -0,0 +1,30 @@
+# REQUIRES: zstd
+# RUN: yaml2obj %s -o %t
+# RUN: obj2yaml %t | FileCheck %s
+
+!Offload
+Compression: zstd
+Members:
+ - ImageKind: IMG_Cubin
+ OffloadKind: OFK_HIP
+ Flags: 0
+ String:
+ - Key: "triple"
+ Value: "amdgcn-amd-amdhsa"
+ - Key: "arch"
+ Value: "gfx90a"
+ Content: "deadbeef"
+
+# CHECK: --- !Offload
+# CHECK-NEXT: Compression: zstd
+# CHECK-NEXT: Members:
+# CHECK-NEXT: - ImageKind: IMG_Cubin
+# CHECK-NEXT: OffloadKind: OFK_HIP
+# CHECK-NEXT: Flags: 0
+# CHECK-NEXT: String:
+# CHECK-NEXT: - Key: triple
+# CHECK-NEXT: Value: amdgcn-amd-amdhsa
+# CHECK-NEXT: - Key: arch
+# CHECK-NEXT: Value: gfx90a
+# CHECK-NEXT: Content: DEADBEEF
+# CHECK-NEXT: ...
diff --git a/llvm/test/ObjectYAML/Offload/malformed-version.yaml b/llvm/test/ObjectYAML/Offload/malformed-version.yaml
index 99383491acce0..c8202daf22791 100644
--- a/llvm/test/ObjectYAML/Offload/malformed-version.yaml
+++ b/llvm/test/ObjectYAML/Offload/malformed-version.yaml
@@ -1,6 +1,6 @@
# RUN: yaml2obj %s | not obj2yaml 2>&1 | FileCheck %s
!Offload
-Version: 3
+Version: 4
Members:
- ImageKind: IMG_Cubin
OffloadKind: OFK_OpenMP
diff --git a/llvm/test/tools/llvm-objdump/Offloading/compressed.test b/llvm/test/tools/llvm-objdump/Offloading/compressed.test
new file mode 100644
index 0000000000000..7316853999208
--- /dev/null
+++ b/llvm/test/tools/llvm-objdump/Offloading/compressed.test
@@ -0,0 +1,41 @@
+# REQUIRES: zstd
+# UNSUPPORTED: system-zos
+
+# RUN: yaml2obj %s --docnum=1 -o %t.bin
+# RUN: llvm-objdump --offloading %t.bin | FileCheck %s --match-full-lines --strict-whitespace --implicit-check-not={{.}}
+
+# RUN: yaml2obj %s --docnum=2 -o %t.o
+# RUN: llvm-objcopy --update-section .llvm.offloading=%t.bin %t.o
+# RUN: llvm-objdump --offloading %t.o | FileCheck %s --check-prefixes=CHECK,ELF --match-full-lines --strict-whitespace --implicit-check-not={{.}}
+
+--- !Offload
+Compression: zstd
+Members:
+ - ImageKind: IMG_Cubin
+ OffloadKind: OFK_HIP
+ String:
+ - Key: triple
+ Value: amdgcn-amd-amdhsa
+ - Key: arch
+ Value: gfx90a
+ Content: DEADBEEF
+
+--- !ELF
+FileHeader:
+ Class: ELFCLASS64
+ Data: ELFDATA2LSB
+ Type: ET_REL
+Sections:
+ - Name: .llvm.offloading
+ Type: SHT_LLVM_OFFLOADING
+ Flags: [ SHF_EXCLUDE ]
+ AddressAlign: 0x0000000000000008
+
+# ELF:{{.*}}file format elf64-unknown
+# ELF-EMPTY:
+# CHECK:OFFLOADING IMAGE [0]:
+# CHECK-NEXT:kind cubin
+# CHECK-NEXT:arch gfx90a
+# CHECK-NEXT:triple amdgcn-amd-amdhsa
+# CHECK-NEXT:producer hip
+# CHECK-NEXT:image size 4 bytes
diff --git a/llvm/test/tools/llvm-offload-binary/compress.test b/llvm/test/tools/llvm-offload-binary/compress.test
new file mode 100644
index 0000000000000..0f19c5d3ebc3a
--- /dev/null
+++ b/llvm/test/tools/llvm-offload-binary/compress.test
@@ -0,0 +1,19 @@
+# REQUIRES: zstd
+# UNSUPPORTED: system-zos
+
+# RUN: llvm-offload-binary -o %t --compress --compression-format=zstd --image=file=%s,arch=abc,triple=x-y-z
+# RUN: llvm-objdump --offloading %t | FileCheck %s
+# RUN: llvm-offload-binary %t --image=file=%t2,arch=abc,triple=x-y-z
+# RUN: diff %s %t2
+
+# RUN: llvm-offload-binary -o %t.u --image=file=%s,arch=abc,triple=x-y-z
+# RUN: cat %t %t.u > %t.concat
+# RUN: llvm-objdump --offloading %t.concat | FileCheck %s --check-prefixes=CHECK,CONCAT
+
+# CHECK: OFFLOADING IMAGE [0]:
+# CHECK-NEXT: kind <none>
+# CHECK-NEXT: arch abc
+# CHECK-NEXT: triple x-y-z
+# CHECK-NEXT: producer none
+# CONCAT: OFFLOADING IMAGE [1]:
+# CONCAT-NEXT: kind <none>
diff --git a/llvm/tools/llvm-objdump/OffloadDump.cpp b/llvm/tools/llvm-objdump/OffloadDump.cpp
index ae3cda778363a..b2d633f28e5f5 100644
--- a/llvm/tools/llvm-objdump/OffloadDump.cpp
+++ b/llvm/tools/llvm-objdump/OffloadDump.cpp
@@ -160,13 +160,14 @@ void llvm::dumpOffloadBundleFatBinary(const ObjectFile &O, StringRef ArchName) {
}
}
-/// Print the contents of an offload binary file \p OB. This may contain
-/// multiple binaries stored in the same buffer.
-void llvm::dumpOffloadSections(const OffloadBinary &OB) {
+/// Print the contents of an offload binary file. This may contain multiple
+/// binaries stored in the same buffer.
+void llvm::dumpOffloadSections(MemoryBufferRef Buffer) {
SmallVector<OffloadFile> Binaries;
- if (Error Err = extractOffloadBinaries(OB.getMemoryBufferRef(), Binaries))
- reportError(OB.getFileName(), "while extracting offloading files: " +
- toString(std::move(Err)));
+ if (Error Err = extractOffloadBinaries(Buffer, Binaries))
+ reportError(Buffer.getBufferIdentifier(),
+ "while extracting offloading files: " +
+ toString(std::move(Err)));
// Print out all the binaries that are contained in this buffer.
for (uint64_t I = 0, E = Binaries.size(); I != E; ++I)
diff --git a/llvm/tools/llvm-objdump/OffloadDump.h b/llvm/tools/llvm-objdump/OffloadDump.h
index 229d479ae357b..eb9a7831f2037 100644
--- a/llvm/tools/llvm-objdump/OffloadDump.h
+++ b/llvm/tools/llvm-objdump/OffloadDump.h
@@ -15,7 +15,7 @@
namespace llvm {
-void dumpOffloadSections(const object::OffloadBinary &OB);
+void dumpOffloadSections(MemoryBufferRef Buffer);
void dumpOffloadBinary(const object::ObjectFile &O, StringRef ArchName);
/// Dump fat binary in binary clang-offload-bundler format
diff --git a/llvm/tools/llvm-objdump/llvm-objdump.cpp b/llvm/tools/llvm-objdump/llvm-objdump.cpp
index 92c1c2da3d025..eaadfd9f1ee45 100644
--- a/llvm/tools/llvm-objdump/llvm-objdump.cpp
+++ b/llvm/tools/llvm-objdump/llvm-objdump.cpp
@@ -3690,9 +3690,10 @@ static void dumpInput(StringRef file) {
dumpObject(O);
else if (MachOUniversalBinary *UB = dyn_cast<MachOUniversalBinary>(&Binary))
parseInputMachO(UB);
- else if (OffloadBinary *OB = dyn_cast<OffloadBinary>(&Binary))
+ else if (isa<OffloadBinary>(&Binary)) {
+ std::unique_ptr<MemoryBuffer> OB = OBinary.takeBinary().second;
dumpOffloadSections(*OB);
- else
+ } else
reportError(errorCodeToError(object_error::invalid_file_type), file);
}
diff --git a/llvm/tools/llvm-offload-binary/llvm-offload-binary.cpp b/llvm/tools/llvm-offload-binary/llvm-offload-binary.cpp
index f3099f43147af..d1dd1f6a34b89 100644
--- a/llvm/tools/llvm-offload-binary/llvm-offload-binary.cpp
+++ b/llvm/tools/llvm-offload-binary/llvm-offload-binary.cpp
@@ -20,6 +20,7 @@
#include "llvm/Object/ObjectFile.h"
#include "llvm/Object/OffloadBinary.h"
#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Compression.h"
#include "llvm/Support/FileOutputBuffer.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
@@ -56,6 +57,22 @@ static cl::opt<bool>
cl::desc("Write extracted files to a static archive"),
cl::cat(OffloadBinaryCategory));
+static cl::opt<bool> Compress("compress",
+ cl::desc("Compress the packaged offload binary"),
+ cl::cat(OffloadBinaryCategory));
+
+static cl::opt<compression::Format> CompressionFormat(
+ "compression-format", cl::desc("Format used with --compress"),
+ cl::values(
+ clEnumValN(compression::Format::Zstd, "zstd", "Zstandard compression"),
+ clEnumValN(compression::Format::Zlib, "zlib", "zlib compression")),
+ cl::init(compression::Format::Zstd), cl::cat(OffloadBinaryCategory));
+
+static cl::opt<int>
+ CompressionLevel("compression-level",
+ cl::desc("Compression level used with --compress"),
+ cl::init(-1), cl::cat(OffloadBinaryCategory));
+
/// Path of the current binary.
static const char *PackagerExecutable;
@@ -126,8 +143,23 @@ static Error bundleImages() {
}
}
- SmallString<0> Buffer = OffloadBinary::write(AllImages);
- if (Buffer.size() % OffloadBinary::getAlignment() != 0)
+ SmallString<0> Buffer;
+ if (Compress) {
+ if (const char *Reason =
+ compression::getReasonIfUnsupported(CompressionFormat))
+ return createStringError(inconvertibleErrorCode(), Reason);
+ compression::Params Params(CompressionFormat);
+ if (CompressionLevel >= 0)
+ Params.level = CompressionLevel;
+ Expected<SmallString<0>> CompressedOrErr =
+ OffloadBinary::write(AllImages, Params);
+ if (!CompressedOrErr)
+ return CompressedOrErr.takeError();
+ Buffer = std::move(*CompressedOrErr);
+ } else {
+ Buffer = OffloadBinary::write(AllImages);
+ }
+ if (!Compress && Buffer.size() % OffloadBinary::getAlignment() != 0)
return createStringError(inconvertibleErrorCode(),
"Offload binary has invalid size alignment");
diff --git a/llvm/tools/obj2yaml/offload2yaml.cpp b/llvm/tools/obj2yaml/offload2yaml.cpp
index 47e86f75514c0..629340128e9a3 100644
--- a/llvm/tools/obj2yaml/offload2yaml.cpp
+++ b/llvm/tools/obj2yaml/offload2yaml.cpp
@@ -10,8 +10,13 @@
#include "llvm/BinaryFormat/Magic.h"
#include "llvm/Object/OffloadBinary.h"
#include "llvm/ObjectYAML/OffloadYAML.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/Support/Compression.h"
+#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/StringSaver.h"
+#include <memory>
+
using namespace llvm;
namespace {
@@ -50,6 +55,25 @@ Expected<OffloadYAML::Binary *> dump(MemoryBufferRef Source,
while (Offset < Source.getBufferSize()) {
MemoryBufferRef Buffer = MemoryBufferRef(
Source.getBuffer().drop_front(Offset), Source.getBufferIdentifier());
+ std::unique_ptr<MemoryBuffer> Aligned;
+ if (!isAddrAligned(Align(object::OffloadBinary::getAlignment()),
+ Buffer.getBufferStart())) {
+ Aligned = MemoryBuffer::getMemBufferCopy(Buffer.getBuffer(),
+ Buffer.getBufferIdentifier());
+ Buffer = *Aligned;
+ }
+ auto HeaderOrErr = object::OffloadBinary::extractHeader(Buffer);
+ if (!HeaderOrErr)
+ return HeaderOrErr.takeError();
+ const object::OffloadBinary::Header *TheHeader = *HeaderOrErr;
+ uint64_t Size = TheHeader->Size;
+ if (TheHeader->Version >= 3 && TheHeader->InflatedSize != 0) {
+ StringRef Payload = Buffer.getBuffer().take_front(Size).drop_front(
+ TheHeader->EntriesOffset);
+ YAMLBinary->Compression = identify_magic(Payload) == file_magic::zstd
+ ? compression::Format::Zstd
+ : compression::Format::Zlib;
+ }
auto BinariesOrErr = object::OffloadBinary::create(Buffer);
if (!BinariesOrErr)
return BinariesOrErr.takeError();
@@ -58,7 +82,7 @@ Expected<OffloadYAML::Binary *> dump(MemoryBufferRef Source,
*BinariesOrErr;
populateYAML(*YAMLBinary, Binaries, Saver);
- Offset += Binaries[0]->getSize();
+ Offset += Size;
}
return YAMLBinary.release();
diff --git a/llvm/unittests/Object/OffloadingTest.cpp b/llvm/unittests/Object/OffloadingTest.cpp
index b6ad6b69f25fc..910ec16cdf58a 100644
--- a/llvm/unittests/Object/OffloadingTest.cpp
+++ b/llvm/unittests/Object/OffloadingTest.cpp
@@ -1,7 +1,13 @@
#include "llvm/Object/OffloadBinary.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/BinaryFormat/Magic.h"
+#include "llvm/Support/Compression.h"
+#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Testing/Support/Error.h"
#include "gtest/gtest.h"
+#include <optional>
#include <random>
using namespace llvm;
@@ -273,3 +279,65 @@ TEST(OffloadingTest, checkEdgeCases) {
EXPECT_EQ(Binaries[0]->getString("large_key").size(), 4096u);
}
}
+
+TEST(OffloadingTest, checkCompressedRoundTrip) {
+ std::optional<compression::Params> Params;
+ if (compression::zstd::isAvailable())
+ Params = compression::Params(compression::Format::Zstd);
+ else if (compression::zlib::isAvailable())
+ Params = compression::Params(compression::Format::Zlib);
+ else
+ GTEST_SKIP() << "compression is unsupported";
+
+ std::string ImageContent(4096, 'A');
+ ImageContent += std::string(4096, 'B');
+
+ OffloadBinary::OffloadingImage Data;
+ Data.TheImageKind = IMG_Object;
+ Data.TheOffloadKind = OFK_HIP;
+ Data.Flags = 7;
+ Data.StringData["triple"] = "amdgcn-amd-amdhsa";
+ Data.StringData["arch"] = "gfx90a";
+ Data.Image = MemoryBuffer::getMemBuffer(ImageContent, "", false);
+
+ SmallString<0> Uncompressed = OffloadBinary::write(Data);
+ Expected<SmallString<0>> CompressedOrErr =
+ OffloadBinary::write(Data, *Params);
+ ASSERT_THAT_EXPECTED(CompressedOrErr, Succeeded());
+ SmallString<0> &Compressed = *CompressedOrErr;
+ EXPECT_LT(Compressed.size(), Uncompressed.size());
+
+ // Verify that the uncompressed and compressed versions are identical.
+ auto HeaderOrErr = OffloadBinary::extractHeader(MemoryBufferRef(
+ StringRef(Compressed.data(), Compressed.size()), "compressed"));
+ ASSERT_THAT_EXPECTED(HeaderOrErr, Succeeded());
+ EXPECT_EQ((*HeaderOrErr)->Size, Compressed.size());
+ EXPECT_EQ((*HeaderOrErr)->InflatedSize, Uncompressed.size());
+
+ auto BinaryBuffer = MemoryBuffer::getMemBufferCopy(Compressed);
+ auto BinariesOrErr = OffloadBinary::create(*BinaryBuffer);
+ ASSERT_THAT_EXPECTED(BinariesOrErr, Succeeded());
+ ASSERT_EQ(BinariesOrErr->size(), 1u);
+
+ OffloadBinary &Binary = *(*BinariesOrErr)[0];
+ EXPECT_EQ(Binary.getImageKind(), IMG_Object);
+ EXPECT_EQ(Binary.getOffloadKind(), OFK_HIP);
+ EXPECT_EQ(Binary.getFlags(), 7u);
+ EXPECT_EQ(Binary.getTriple(), "amdgcn-amd-amdhsa");
+ EXPECT_EQ(Binary.getArch(), "gfx90a");
+ EXPECT_EQ(Binary.getImage(), ImageContent);
+ EXPECT_EQ(Binary.getSize(), Uncompressed.size());
+
+ // The concatenated form should still extract properly when compressed.
+ SmallString<0> Concat = Compressed;
+ Concat.append(Uncompressed);
+ SmallVector<OffloadFile> Files;
+ ASSERT_THAT_ERROR(
+ extractOffloadBinaries(
+ MemoryBufferRef(StringRef(Concat.data(), Concat.size()), "concat"),
+ Files),
+ Succeeded());
+ ASSERT_EQ(Files.size(), 2u);
+ EXPECT_EQ(Files[0].getBinary()->getImage(), ImageContent);
+ EXPECT_EQ(Files[1].getBinary()->getImage(), ImageContent);
+}
More information about the llvm-branch-commits
mailing list