[llvm] [llvm-objdump] Add support for HIP offload bundles (PR #114834)
Joseph Huber via llvm-commits
llvm-commits at lists.llvm.org
Wed Mar 19 12:52:42 PDT 2025
================
@@ -0,0 +1,477 @@
+//===- OffloadBundle.cpp - Utilities for handling offloading code -*- C++
+//-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Object/OffloadBundle.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/BinaryFormat/Magic.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/MC/StringTableBuilder.h"
+#include "llvm/Object/Archive.h"
+#include "llvm/Object/Binary.h"
+#include "llvm/Object/COFF.h"
+#include "llvm/Object/ELFObjectFile.h"
+#include "llvm/Object/Error.h"
+#include "llvm/Object/IRObjectFile.h"
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/Support/BinaryStreamReader.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/Timer.h"
+
+using namespace llvm;
+using namespace llvm::object;
+
+static llvm::TimerGroup
+ OffloadBundlerTimerGroup("Offload Bundler Timer Group",
+ "Timer group for offload bundler");
+
+// Extract an Offload bundle (usually a Offload Bundle) from a fat_bin
+// section
+Error extractOffloadBundle(MemoryBufferRef Contents, uint64_t SectionOffset,
+ StringRef FileName,
+ SmallVectorImpl<OffloadBundleFatBin> &Bundles) {
+
+ uint64_t Offset = 0;
+ int64_t NextbundleStart = 0;
+
+ // There could be multiple offloading bundles stored at this section.
+ while (NextbundleStart >= 0) {
+
+ std::unique_ptr<MemoryBuffer> Buffer =
+ MemoryBuffer::getMemBuffer(Contents.getBuffer().drop_front(Offset), "",
+ /*RequiresNullTerminator=*/false);
+
+ // Create the FatBinBindle object. This will also create the Bundle Entry
+ // list info.
+ auto FatBundleOrErr =
+ OffloadBundleFatBin::create(*Buffer, SectionOffset + Offset, FileName);
+ if (!FatBundleOrErr)
+ return FatBundleOrErr.takeError();
+
+ // Add current Bundle to list.
+ Bundles.emplace_back(std::move(**FatBundleOrErr));
+
+ // Find the next bundle by searching for the magic string
+ StringRef Str = Buffer->getBuffer();
+ NextbundleStart =
+ (int64_t)Str.find(StringRef("__CLANG_OFFLOAD_BUNDLE__"), 24);
+
+ if (NextbundleStart >= 0)
+ Offset += NextbundleStart;
+ }
+
+ return Error::success();
+}
+
+Error OffloadBundleFatBin::readEntries(StringRef Buffer,
+ uint64_t SectionOffset) {
+ uint64_t NumOfEntries = 0;
+
+ BinaryStreamReader Reader(Buffer, llvm::endianness::little);
+
+ // Read the Magic String first.
+ StringRef Magic;
+ if (auto EC = Reader.readFixedString(Magic, 24)) {
+ return errorCodeToError(object_error::parse_failed);
+ }
+
+ // Read the number of Code Objects (Entries) in the current Bundle.
+ if (auto EC = Reader.readInteger(NumOfEntries))
+ return errorCodeToError(object_error::parse_failed);
+
+ NumberOfEntries = NumOfEntries;
+
+ // For each Bundle Entry (code object)
+ for (uint64_t I = 0; I < NumOfEntries; I++) {
+ uint64_t EntrySize;
+ uint64_t EntryOffset;
+ uint64_t EntryIDSize;
+ StringRef EntryID;
+
+ if (auto EC = Reader.readInteger(EntryOffset))
+ return errorCodeToError(object_error::parse_failed);
+
+ if (auto EC = Reader.readInteger(EntrySize))
+ return errorCodeToError(object_error::parse_failed);
+
+ if (auto EC = Reader.readInteger(EntryIDSize))
+ return errorCodeToError(object_error::parse_failed);
+
+ if (auto EC = Reader.readFixedString(EntryID, EntryIDSize))
+ return errorCodeToError(object_error::parse_failed);
+
+ // create a Bundle Entry object:
+ auto Entry = new OffloadBundleEntry(EntryOffset + SectionOffset, EntrySize,
+ EntryIDSize, EntryID);
+
+ Entries.push_back(*Entry);
+ }
+
+ return Error::success();
+}
+
+Expected<std::unique_ptr<OffloadBundleFatBin>>
+OffloadBundleFatBin::create(MemoryBufferRef Buf, uint64_t SectionOffset,
+ StringRef FileName) {
+ if (Buf.getBufferSize() < 24)
+ return errorCodeToError(object_error::parse_failed);
+
+ // Check for magic bytes.
+ if (identify_magic(Buf.getBuffer()) != file_magic::offload_bundle)
+ return errorCodeToError(object_error::parse_failed);
+
+ OffloadBundleFatBin *TheBundle = new OffloadBundleFatBin(Buf, FileName);
+
+ // Read the Bundle Entries
+ Error Err = TheBundle->readEntries(Buf.getBuffer(), SectionOffset);
+ if (Err)
+ return errorCodeToError(object_error::parse_failed);
+
+ return std::unique_ptr<OffloadBundleFatBin>(TheBundle);
+}
+
+Error OffloadBundleFatBin::extractBundle(const ObjectFile &Source) {
+ // This will extract all entries in the Bundle
+ for (OffloadBundleEntry &Entry : Entries) {
+
+ if (Entry.Size > 0) {
+ // create output file name. Which should be
+ // <fileName>-offset<Offset>-size<Size>.co"
+ std::string Str = getFileName().str() + "-offset" + itostr(Entry.Offset) +
+ "-size" + itostr(Entry.Size) + ".co";
+ if (Error Err = object::extractCodeObject(Source, Entry.Offset,
+ Entry.Size, StringRef(Str)))
+ return Err;
+ }
+ }
+
+ return Error::success();
+}
+
+Error object::extractOffloadBundleFatBinary(
+ const ObjectFile &Obj, SmallVectorImpl<OffloadBundleFatBin> &Bundles) {
+ assert((Obj.isELF() || Obj.isCOFF()) && "Invalid file type");
+
+ // iterate through Sections until we find an offload_bundle section.
+ for (SectionRef Sec : Obj.sections()) {
+ Expected<StringRef> Buffer = Sec.getContents();
+ if (!Buffer)
+ return Buffer.takeError();
+
+ // If it does not start with the reserved suffix, just skip this section.
+ if ((llvm::identify_magic(*Buffer) == llvm::file_magic::offload_bundle) ||
+ (llvm::identify_magic(*Buffer) ==
+ llvm::file_magic::offload_bundle_compressed)) {
+
+ uint64_t SectionOffset = 0;
+ if (Obj.isELF()) {
+ SectionOffset = ELFSectionRef(Sec).getOffset();
+ } else if (Obj.isCOFF()) {
+ if (const COFFObjectFile *COFFObj = dyn_cast<COFFObjectFile>(&Obj)) {
+ const coff_section *CoffSection = COFFObj->getCOFFSection(Sec);
+ }
+ }
+
+ MemoryBufferRef Contents(*Buffer, Obj.getFileName());
+
+ if (llvm::identify_magic(*Buffer) ==
+ llvm::file_magic::offload_bundle_compressed) {
+ // Decompress the input if necessary.
+ Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =
+ CompressedOffloadBundle::decompress(Contents, false);
+
+ if (!DecompressedBufferOrErr)
+ return createStringError(
+ inconvertibleErrorCode(),
+ "Failed to decompress input: " +
+ llvm::toString(DecompressedBufferOrErr.takeError()));
+
+ MemoryBuffer &DecompressedInput = **DecompressedBufferOrErr;
+ if (Error Err = extractOffloadBundle(DecompressedInput, SectionOffset,
+ Obj.getFileName(), Bundles))
+ return Err;
+ } else {
+ if (Error Err = extractOffloadBundle(Contents, SectionOffset,
+ Obj.getFileName(), Bundles))
+ return Err;
+ }
+ }
+ }
+ return Error::success();
+}
+
+Error object::extractCodeObject(const ObjectFile &Source, int64_t Offset,
+ int64_t Size, StringRef OutputFileName) {
+ Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
+ FileOutputBuffer::create(OutputFileName, Size);
+
+ if (!BufferOrErr)
+ return BufferOrErr.takeError();
+
+ Expected<MemoryBufferRef> InputBuffOrErr = Source.getMemoryBufferRef();
+ if (Error Err = InputBuffOrErr.takeError())
+ return Err;
+
+ std::unique_ptr<FileOutputBuffer> Buf = std::move(*BufferOrErr);
+ std::copy(InputBuffOrErr->getBufferStart() + Offset,
+ InputBuffOrErr->getBufferStart() + Offset + Size,
+ Buf->getBufferStart());
+ if (Error E = Buf->commit())
+ return E;
+
+ return Error::success();
+}
+
+// given a file name, offset, and size, extract data into a code object file,
+// into file <SourceFile>-offset<Offset>-size<Size>.co
+Error object::extractOffloadBundleByURI(StringRef URIstr) {
+ // create a URI object
+ Expected<std::unique_ptr<OffloadBundleURI>> UriOrErr(
+ OffloadBundleURI::createOffloadBundleURI(URIstr, FILE_URI));
+ if (!UriOrErr)
+ return UriOrErr.takeError();
+
+ OffloadBundleURI &Uri = **UriOrErr;
+ std::string OutputFile = Uri.FileName.str();
+ OutputFile +=
+ "-offset" + itostr(Uri.Offset) + "-size" + itostr(Uri.Size) + ".co";
+
+ // Create an ObjectFile object from uri.file_uri
+ auto ObjOrErr = ObjectFile::createObjectFile(Uri.FileName);
+ if (!ObjOrErr)
+ return ObjOrErr.takeError();
+
+ auto Obj = ObjOrErr->getBinary();
+ if (Error Err =
+ object::extractCodeObject(*Obj, Uri.Offset, Uri.Size, OutputFile))
+ return Err;
+
+ return Error::success();
+}
+
+// Utility function to format numbers with commas
+static std::string formatWithCommas(unsigned long long Value) {
+ std::string Num = std::to_string(Value);
+ int InsertPosition = Num.length() - 3;
+ while (InsertPosition > 0) {
+ Num.insert(InsertPosition, ",");
+ InsertPosition -= 3;
+ }
+ return Num;
+}
+
+llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
+CompressedOffloadBundle::decompress(llvm::MemoryBufferRef &Input,
+ bool Verbose) {
+ StringRef Blob = Input.getBuffer();
+
+ if (Blob.size() < V1HeaderSize)
+ return llvm::MemoryBuffer::getMemBufferCopy(Blob);
+
+ if (llvm::identify_magic(Blob) !=
+ llvm::file_magic::offload_bundle_compressed) {
+ if (Verbose)
+ llvm::errs() << "Uncompressed bundle.\n";
+ return llvm::MemoryBuffer::getMemBufferCopy(Blob);
+ }
+
+ size_t CurrentOffset = MagicSize;
+
+ 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 TotalFileSize;
+ if (ThisVersion >= 2) {
+ if (Blob.size() < V2HeaderSize)
+ return createStringError(inconvertibleErrorCode(),
+ "Compressed bundle header size too small");
+ memcpy(&TotalFileSize, Blob.data() + CurrentOffset, sizeof(uint32_t));
+ CurrentOffset += FileSizeFieldSize;
+ }
+
+ uint32_t UncompressedSize;
+ memcpy(&UncompressedSize, Blob.data() + CurrentOffset, sizeof(uint32_t));
+ CurrentOffset += UncompressedSizeFieldSize;
+
+ uint64_t StoredHash;
+ memcpy(&StoredHash, Blob.data() + CurrentOffset, sizeof(uint64_t));
+ CurrentOffset += HashFieldSize;
+
+ llvm::compression::Format CompressionFormat;
+ if (CompressionMethod ==
+ static_cast<uint16_t>(llvm::compression::Format::Zlib))
+ CompressionFormat = llvm::compression::Format::Zlib;
+ else if (CompressionMethod ==
+ static_cast<uint16_t>(llvm::compression::Format::Zstd))
+ CompressionFormat = llvm::compression::Format::Zstd;
+ else
+ return createStringError(inconvertibleErrorCode(),
+ "Unknown compressing method");
+
+ llvm::Timer DecompressTimer("Decompression Timer", "Decompression time",
+ OffloadBundlerTimerGroup);
+ if (Verbose)
+ DecompressTimer.startTimer();
+
+ SmallVector<uint8_t, 0> DecompressedData;
+ StringRef CompressedData = Blob.substr(CurrentOffset);
+ if (llvm::Error DecompressionError = llvm::compression::decompress(
+ CompressionFormat, llvm::arrayRefFromStringRef(CompressedData),
+ DecompressedData, UncompressedSize))
+ return createStringError(inconvertibleErrorCode(),
+ "Could not decompress embedded file contents: " +
+ llvm::toString(std::move(DecompressionError)));
+
+ if (Verbose) {
+ DecompressTimer.stopTimer();
+
+ double DecompressionTimeSeconds =
+ DecompressTimer.getTotalTime().getWallTime();
+
+ // Recalculate MD5 hash for integrity check
----------------
jhuber6 wrote:
```suggestion
// Recalculate MD5 hash for integrity check.
```
https://github.com/llvm/llvm-project/pull/114834
More information about the llvm-commits
mailing list