https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/222774
>From a73ef9fc74e7e33a3e74a99dd15a7c4a96158109 Mon Sep 17 00:00:00 2001 From: Joseph Huber <[email protected]> Date: Fri, 18 Sep 2026 11:53:57 -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/docs/CommandGuide/llvm-offload-binary.md | 29 ++- llvm/docs/ReleaseNotes.md | 2 + llvm/include/llvm/Object/OffloadBinary.h | 47 +++-- llvm/include/llvm/ObjectYAML/OffloadYAML.h | 5 + llvm/lib/Object/OffloadBinary.cpp | 183 +++++++++++++----- llvm/lib/ObjectYAML/OffloadEmitter.cpp | 15 +- llvm/lib/ObjectYAML/OffloadYAML.cpp | 7 + .../ObjectYAML/Offload/compressed-zlib.yaml | 29 +++ llvm/test/ObjectYAML/Offload/compressed.yaml | 29 +++ .../Offload/malformed-compressed-offset.yaml | 18 ++ .../ObjectYAML/Offload/malformed-version.yaml | 2 +- .../llvm-objdump/Offloading/compressed.test | 41 ++++ .../tools/llvm-offload-binary/compress.test | 23 +++ 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 | 38 +++- llvm/tools/obj2yaml/offload2yaml.cpp | 31 ++- llvm/unittests/Object/OffloadingTest.cpp | 78 ++++++++ 19 files changed, 500 insertions(+), 97 deletions(-) create mode 100644 llvm/test/ObjectYAML/Offload/compressed-zlib.yaml create mode 100644 llvm/test/ObjectYAML/Offload/compressed.yaml create mode 100644 llvm/test/ObjectYAML/Offload/malformed-compressed-offset.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/docs/CommandGuide/llvm-offload-binary.md b/llvm/docs/CommandGuide/llvm-offload-binary.md index b3ac0386ebe1b..6af4bbe7d4ffa 100644 --- a/llvm/docs/CommandGuide/llvm-offload-binary.md +++ b/llvm/docs/CommandGuide/llvm-offload-binary.md @@ -29,6 +29,8 @@ version and size. Each binary contains its own header, allowing tools to locate offloading sections even when merged by a linker. Each offload entry includes metadata such as the device image kind, producer kind, and key-value string metadata. Multiple offloading images are concatenated to form a fat binary. +Images can be compressed with zlib or zstd for better size efficiency. +Extraction handles compressed binaries automatically. ## EXAMPLE @@ -62,6 +64,18 @@ When extracting from an input binary, write all extracted images into a static archive instead of separate files. ::: +:::{option} --compress +Compress the packaged offload binary. +::: + +:::{option} --compression-format=<zstd|zlib> +Select the compression format. The default is `zstd`. +::: + +:::{option} --compression-level=<level> +Specify a compression level for the selected format. +::: + :::{option} --image=<<key>=<value>,...> Specify a set of arbitrary key-value arguments describing an image. Commonly used optional keys include `arch` (e.g. `sm_70` for CUDA) and @@ -112,13 +126,14 @@ an image buffer. :name: table-binary_header :widths: auto -| Type | Identifier | Description | -| -------- | ------------ | -------------------------------------------------- | -| uint8_t | magic | The magic bytes for the binary format (0x10FF10AD) | -| uint32_t | version | Version of this format (currently version 1) | -| uint64_t | size | Size of this binary in bytes | -| uint64_t | entry offset | Absolute offset of the offload entries in bytes | -| uint64_t | entry size | Size of the offload entries in bytes | +| Type | Identifier | Description | +| -------- | ------------- | --------------------------------------------------- | +| uint32_t | magic | The magic bytes for the binary format (0x10FF10AD) | +| uint32_t | version | Version of this format (currently version 3) | +| uint64_t | size | Size of this binary in bytes | +| uint64_t | entry offset | Absolute offset of the offload entries in bytes | +| uint64_t | entry count | Number of offload entries | +| uint64_t | inflated size | Uncompressed size, or zero if compression is unused | ::: Each offload entry describes a bundled image along with its associated metadata. diff --git a/llvm/docs/ReleaseNotes.md b/llvm/docs/ReleaseNotes.md index 985117034e2be..e7244121ec532 100644 --- a/llvm/docs/ReleaseNotes.md +++ b/llvm/docs/ReleaseNotes.md @@ -295,6 +295,8 @@ Makes programs 10x faster by doing Special New Thing. ### Changes to the LLVM tools +* llvm-offload-binary can now compress packaged binaries using zstd or zlib. + * llvm-mca no longer defaults -mcpu to "native" * llvm-rc now supports `/showIncludes` to report header and resource-file diff --git a/llvm/include/llvm/Object/OffloadBinary.h b/llvm/include/llvm/Object/OffloadBinary.h index ab664800567e4..0c4d5cb1739e7 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. @@ -199,30 +209,26 @@ class OffloadBinary : public Binary { const uint64_t Index; }; -/// A class to contain the binary information for a single OffloadBinary. -/// Memory is shared between multiple OffloadBinary instances read from -/// the single serialized offload binary. -class OffloadFile : public OwningBinary<OffloadBinary> { +/// A handle to a single OffloadBinary. +class OffloadFile { public: using TargetID = std::pair<StringRef, StringRef>; - OffloadFile(std::unique_ptr<OffloadBinary> Binary, - std::unique_ptr<MemoryBuffer> Buffer) - : OwningBinary<OffloadBinary>(std::move(Binary), std::move(Buffer)) {} + OffloadFile() = default; + OffloadFile(std::unique_ptr<OffloadBinary> Binary) + : Binary(std::move(Binary)) {} + + OffloadBinary *getBinary() { return Binary.get(); } + const OffloadBinary *getBinary() const { return Binary.get(); } /// Make a deep copy of this offloading file. OffloadFile copy() const { - std::unique_ptr<MemoryBuffer> Buffer = MemoryBuffer::getMemBufferCopy( - getBinary()->getMemoryBufferRef().getBuffer(), - getBinary()->getMemoryBufferRef().getBufferIdentifier()); - - // This parsing should never fail because it has already been parsed. - auto NewBinaryOrErr = - OffloadBinary::create(*Buffer, getBinary()->getIndex()); + auto NewBinaryOrErr = OffloadBinary::create( + getBinary()->getMemoryBufferRef(), getBinary()->getIndex()); assert(NewBinaryOrErr && "Failed to parse a copy of the binary?"); if (!NewBinaryOrErr) llvm::consumeError(NewBinaryOrErr.takeError()); - return OffloadFile(std::move((*NewBinaryOrErr)[0]), std::move(Buffer)); + return OffloadFile(std::move((*NewBinaryOrErr)[0])); } /// We use the Triple and Architecture pair to group linker inputs together. @@ -230,6 +236,9 @@ class OffloadFile : public OwningBinary<OffloadBinary> { operator TargetID() const { return std::make_pair(getBinary()->getTriple(), getBinary()->getArch()); } + +private: + std::unique_ptr<OffloadBinary> Binary; }; /// Extracts embedded device offloading code from a memory \p Buffer to a list 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..89ffd623612ba 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,26 +59,16 @@ 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)); - } + for (auto &Binary : *BinariesOrErr) + Binaries.emplace_back(std::move(Binary)); - Offset += Header->Size; + Offset = alignTo(Offset + Header->Size, OffloadBinary::getAlignment()); } return Error::success(); @@ -197,11 +176,53 @@ 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 != sizeof(OffloadBinary::Header) || + 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; + + ArrayRef<uint8_t> CompressedBytes = arrayRefFromStringRef(Compressed); + if (const char *Reason = compression::getReasonIfUnsupported(CompressedBytes)) + return createStringError(Reason); + + SmallVector<uint8_t, 0> Body; + if (Error Err = compression::decompress(CompressedBytes, 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), + sizeof(OffloadBinary::Header))); + 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 +238,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 +260,45 @@ 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() || - TheEntry->StringOffset + TheEntry->NumStrings * sizeof(StringEntry) > - Buf.getBufferSize()) + const uint64_t BufSize = Owned.getBufferSize(); + if (TheEntry->ImageOffset > BufSize || + TheEntry->ImageSize > BufSize - TheEntry->ImageOffset) + return errorCodeToError(object_error::unexpected_eof); + + const size_t StringEntrySize = + TheHeader->Version == 1 ? sizeof(StringEntryV1) : sizeof(StringEntry); + if (TheEntry->StringOffset > BufSize || + TheEntry->NumStrings > + (BufSize - TheEntry->StringOffset) / StringEntrySize) return errorCodeToError(object_error::unexpected_eof); return Error::success(); }; @@ -259,7 +311,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 +322,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 +358,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 +416,38 @@ 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(alignTo(Header.Size, getAlignment())); + raw_svector_ostream OS(Data); + OS << StringRef(reinterpret_cast<const char *>(&Header), + Header.EntriesOffset); + OS << toStringRef(CompressedBuffer); + assert(Header.Size == OS.tell() && "Size mismatch"); + OS.write_zeros(alignTo(Header.Size, getAlignment()) - Header.Size); + 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-zlib.yaml b/llvm/test/ObjectYAML/Offload/compressed-zlib.yaml new file mode 100644 index 0000000000000..4be198ef3aecb --- /dev/null +++ b/llvm/test/ObjectYAML/Offload/compressed-zlib.yaml @@ -0,0 +1,29 @@ +# REQUIRES: zlib +# RUN: yaml2obj %s -o %t +# RUN: obj2yaml %t | FileCheck %s + +!Offload +Compression: zlib +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: 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/compressed.yaml b/llvm/test/ObjectYAML/Offload/compressed.yaml new file mode 100644 index 0000000000000..3af5f6c0e7854 --- /dev/null +++ b/llvm/test/ObjectYAML/Offload/compressed.yaml @@ -0,0 +1,29 @@ +# 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: 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-compressed-offset.yaml b/llvm/test/ObjectYAML/Offload/malformed-compressed-offset.yaml new file mode 100644 index 0000000000000..02210266aa141 --- /dev/null +++ b/llvm/test/ObjectYAML/Offload/malformed-compressed-offset.yaml @@ -0,0 +1,18 @@ +# REQUIRES: zstd +# RUN: yaml2obj %s | not obj2yaml 2>&1 | FileCheck %s + +!Offload +Compression: zstd +EntriesOffset: 48 +Members: + - ImageKind: IMG_Cubin + OffloadKind: OFK_HIP + Flags: 0 + String: + - Key: "triple" + Value: "amdgcn-amd-amdhsa" + - Key: "arch" + Value: "gfx90a" + Content: "deadbeef" + +# CHECK: Error reading file: <stdin>: The end of the file was unexpectedly encountered 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..abaed92855f74 --- /dev/null +++ b/llvm/test/tools/llvm-offload-binary/compress.test @@ -0,0 +1,23 @@ +# 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.neg --compress --compression-format=zstd --compression-level=-5 --image=file=%s,arch=abc,triple=x-y-z +# RUN: llvm-offload-binary %t.neg --image=file=%t.neg.out,arch=abc,triple=x-y-z +# RUN: diff %s %t.neg.out + +# 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 2d670fd338309..12fadcdae7326 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..80124b962753e 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; @@ -88,9 +105,9 @@ static Error writeFile(StringRef Filename, StringRef Data) { static Error bundleImages() { SmallVector<OffloadBinary::OffloadingImage> AllImages; + BumpPtrAllocator Alloc; + StringSaver Saver(Alloc); for (StringRef Image : DeviceImages) { - BumpPtrAllocator Alloc; - StringSaver Saver(Alloc); DenseMap<StringRef, StringRef> Args = getImageArguments(Image, Saver); if (!Args.count("file")) @@ -126,7 +143,22 @@ static Error bundleImages() { } } - SmallString<0> Buffer = OffloadBinary::write(AllImages); + SmallString<0> Buffer; + if (Compress) { + if (const char *Reason = + compression::getReasonIfUnsupported(CompressionFormat)) + return createStringError(inconvertibleErrorCode(), Reason); + compression::Params Params(CompressionFormat); + if (CompressionLevel.getNumOccurrences()) + 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 (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..d7c1f0f8ceb93 100644 --- a/llvm/tools/obj2yaml/offload2yaml.cpp +++ b/llvm/tools/obj2yaml/offload2yaml.cpp @@ -7,18 +7,21 @@ //===----------------------------------------------------------------------===// #include "obj2yaml.h" -#include "llvm/BinaryFormat/Magic.h" #include "llvm/Object/OffloadBinary.h" #include "llvm/ObjectYAML/OffloadYAML.h" +#include "llvm/Support/Alignment.h" +#include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/StringSaver.h" +#include <memory> + using namespace llvm; namespace { void populateYAML(OffloadYAML::Binary &YAMLBinary, ArrayRef<std::unique_ptr<object::OffloadBinary>> OBinaries, - UniqueStringSaver Saver) { + UniqueStringSaver &Saver) { for (const auto &OBinaryPtr : OBinaries) { object::OffloadBinary &OB = *OBinaryPtr; @@ -35,12 +38,12 @@ void populateYAML(OffloadYAML::Binary &YAMLBinary, } if (!OB.getImage().empty()) - Member.Content = arrayRefFromStringRef(OB.getImage()); + Member.Content = arrayRefFromStringRef(Saver.save(OB.getImage())); } } Expected<OffloadYAML::Binary *> dump(MemoryBufferRef Source, - UniqueStringSaver Saver) { + UniqueStringSaver &Saver) { std::unique_ptr<OffloadYAML::Binary> YAMLBinary = std::make_unique<OffloadYAML::Binary>(); @@ -50,15 +53,27 @@ 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; + auto BinariesOrErr = object::OffloadBinary::create(Buffer); if (!BinariesOrErr) return BinariesOrErr.takeError(); - SmallVector<std::unique_ptr<object::OffloadBinary>> &Binaries = - *BinariesOrErr; - populateYAML(*YAMLBinary, Binaries, Saver); + populateYAML(*YAMLBinary, *BinariesOrErr, Saver); - Offset += Binaries[0]->getSize(); + Offset = + alignTo(Offset + Size, Align(object::OffloadBinary::getAlignment())); } return YAMLBinary.release(); diff --git a/llvm/unittests/Object/OffloadingTest.cpp b/llvm/unittests/Object/OffloadingTest.cpp index b6ad6b69f25fc..8661bdefa725b 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/Support/Compression.h" +#include "llvm/Support/MathExtras.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,75 @@ 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(Compressed.size(), + alignTo((*HeaderOrErr)->Size, OffloadBinary::getAlignment())); + 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()); + + // Two compressed binaries should extract across the first binary's padding. + OffloadBinary::OffloadingImage SecondData; + SecondData.TheImageKind = IMG_Bitcode; + SecondData.TheOffloadKind = OFK_OpenMP; + SecondData.Image = MemoryBuffer::getMemBuffer( + "SecondImage", "", /*RequiresNullTerminator=*/false); + Expected<SmallString<0>> SecondCompressedOrErr = + OffloadBinary::write(SecondData, *Params); + ASSERT_THAT_EXPECTED(SecondCompressedOrErr, Succeeded()); + + SmallString<0> Concat = Compressed; + Concat.append(*SecondCompressedOrErr); + 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(), "SecondImage"); +} _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
