601 lines · cpp
1//===- OffloadBundle.cpp - Utilities for offload bundles---*- C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------===//8 9#include "llvm/Object/OffloadBundle.h"10#include "llvm/BinaryFormat/Magic.h"11#include "llvm/IR/Module.h"12#include "llvm/IRReader/IRReader.h"13#include "llvm/MC/StringTableBuilder.h"14#include "llvm/Object/Archive.h"15#include "llvm/Object/Binary.h"16#include "llvm/Object/COFF.h"17#include "llvm/Object/ELFObjectFile.h"18#include "llvm/Object/Error.h"19#include "llvm/Object/IRObjectFile.h"20#include "llvm/Object/ObjectFile.h"21#include "llvm/Support/BinaryStreamReader.h"22#include "llvm/Support/SourceMgr.h"23#include "llvm/Support/Timer.h"24 25using namespace llvm;26using namespace llvm::object;27 28static TimerGroup OffloadBundlerTimerGroup("Offload Bundler Timer Group",29 "Timer group for offload bundler");30 31// Extract an Offload bundle (usually a Offload Bundle) from a fat_bin32// section.33Error extractOffloadBundle(MemoryBufferRef Contents, uint64_t SectionOffset,34 StringRef FileName,35 SmallVectorImpl<OffloadBundleFatBin> &Bundles) {36 37 size_t Offset = 0;38 size_t NextbundleStart = 0;39 StringRef Magic;40 std::unique_ptr<MemoryBuffer> Buffer;41 42 // There could be multiple offloading bundles stored at this section.43 while ((NextbundleStart != StringRef::npos) &&44 (Offset < Contents.getBuffer().size())) {45 Buffer =46 MemoryBuffer::getMemBuffer(Contents.getBuffer().drop_front(Offset), "",47 /*RequiresNullTerminator=*/false);48 49 if (identify_magic((*Buffer).getBuffer()) ==50 file_magic::offload_bundle_compressed) {51 Magic = "CCOB";52 // Decompress this bundle first.53 NextbundleStart = (*Buffer).getBuffer().find(Magic, Magic.size());54 if (NextbundleStart == StringRef::npos)55 NextbundleStart = (*Buffer).getBuffer().size();56 57 ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =58 MemoryBuffer::getMemBuffer(59 (*Buffer).getBuffer().take_front(NextbundleStart), FileName,60 false);61 if (std::error_code EC = CodeOrErr.getError())62 return createFileError(FileName, EC);63 64 Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =65 CompressedOffloadBundle::decompress(**CodeOrErr, nullptr);66 if (!DecompressedBufferOrErr)67 return createStringError("failed to decompress input: " +68 toString(DecompressedBufferOrErr.takeError()));69 70 auto FatBundleOrErr = OffloadBundleFatBin::create(71 **DecompressedBufferOrErr, Offset, FileName, true);72 if (!FatBundleOrErr)73 return FatBundleOrErr.takeError();74 75 // Add current Bundle to list.76 Bundles.emplace_back(std::move(**FatBundleOrErr));77 78 } else if (identify_magic((*Buffer).getBuffer()) ==79 file_magic::offload_bundle) {80 // Create the OffloadBundleFatBin object. This will also create the Bundle81 // Entry list info.82 auto FatBundleOrErr = OffloadBundleFatBin::create(83 *Buffer, SectionOffset + Offset, FileName);84 if (!FatBundleOrErr)85 return FatBundleOrErr.takeError();86 87 // Add current Bundle to list.88 Bundles.emplace_back(std::move(**FatBundleOrErr));89 90 Magic = "__CLANG_OFFLOAD_BUNDLE__";91 NextbundleStart = (*Buffer).getBuffer().find(Magic, Magic.size());92 }93 94 if (NextbundleStart != StringRef::npos)95 Offset += NextbundleStart;96 }97 98 return Error::success();99}100 101Error OffloadBundleFatBin::readEntries(StringRef Buffer,102 uint64_t SectionOffset) {103 uint64_t NumOfEntries = 0;104 105 BinaryStreamReader Reader(Buffer, llvm::endianness::little);106 107 // Read the Magic String first.108 StringRef Magic;109 if (auto EC = Reader.readFixedString(Magic, 24))110 return errorCodeToError(object_error::parse_failed);111 112 // Read the number of Code Objects (Entries) in the current Bundle.113 if (auto EC = Reader.readInteger(NumOfEntries))114 return errorCodeToError(object_error::parse_failed);115 116 NumberOfEntries = NumOfEntries;117 118 // For each Bundle Entry (code object).119 for (uint64_t I = 0; I < NumOfEntries; I++) {120 uint64_t EntrySize;121 uint64_t EntryOffset;122 uint64_t EntryIDSize;123 StringRef EntryID;124 125 if (Error Err = Reader.readInteger(EntryOffset))126 return Err;127 128 if (Error Err = Reader.readInteger(EntrySize))129 return Err;130 131 if (Error Err = Reader.readInteger(EntryIDSize))132 return Err;133 134 if (Error Err = Reader.readFixedString(EntryID, EntryIDSize))135 return Err;136 137 auto Entry = std::make_unique<OffloadBundleEntry>(138 EntryOffset + SectionOffset, EntrySize, EntryIDSize, EntryID);139 140 Entries.push_back(*Entry);141 }142 143 return Error::success();144}145 146Expected<std::unique_ptr<OffloadBundleFatBin>>147OffloadBundleFatBin::create(MemoryBufferRef Buf, uint64_t SectionOffset,148 StringRef FileName, bool Decompress) {149 if (Buf.getBufferSize() < 24)150 return errorCodeToError(object_error::parse_failed);151 152 // Check for magic bytes.153 if ((identify_magic(Buf.getBuffer()) != file_magic::offload_bundle) &&154 (identify_magic(Buf.getBuffer()) !=155 file_magic::offload_bundle_compressed))156 return errorCodeToError(object_error::parse_failed);157 158 std::unique_ptr<OffloadBundleFatBin> TheBundle(159 new OffloadBundleFatBin(Buf, FileName));160 161 // Read the Bundle Entries.162 Error Err =163 TheBundle->readEntries(Buf.getBuffer(), Decompress ? 0 : SectionOffset);164 if (Err)165 return Err;166 167 return std::move(TheBundle);168}169 170Error OffloadBundleFatBin::extractBundle(const ObjectFile &Source) {171 // This will extract all entries in the Bundle.172 for (OffloadBundleEntry &Entry : Entries) {173 174 if (Entry.Size == 0)175 continue;176 177 // create output file name. Which should be178 // <fileName>-offset<Offset>-size<Size>.co"179 std::string Str = getFileName().str() + "-offset" + itostr(Entry.Offset) +180 "-size" + itostr(Entry.Size) + ".co";181 if (Error Err = object::extractCodeObject(Source, Entry.Offset, Entry.Size,182 StringRef(Str)))183 return Err;184 }185 186 return Error::success();187}188 189Error object::extractOffloadBundleFatBinary(190 const ObjectFile &Obj, SmallVectorImpl<OffloadBundleFatBin> &Bundles) {191 assert((Obj.isELF() || Obj.isCOFF()) && "Invalid file type");192 193 // Iterate through Sections until we find an offload_bundle section.194 for (SectionRef Sec : Obj.sections()) {195 Expected<StringRef> Buffer = Sec.getContents();196 if (!Buffer)197 return Buffer.takeError();198 199 // If it does not start with the reserved suffix, just skip this section.200 if ((llvm::identify_magic(*Buffer) == file_magic::offload_bundle) ||201 (llvm::identify_magic(*Buffer) ==202 file_magic::offload_bundle_compressed)) {203 204 uint64_t SectionOffset = 0;205 if (Obj.isELF()) {206 SectionOffset = ELFSectionRef(Sec).getOffset();207 } else if (Obj.isCOFF()) // TODO: add COFF Support.208 return createStringError(object_error::parse_failed,209 "COFF object files not supported");210 211 MemoryBufferRef Contents(*Buffer, Obj.getFileName());212 if (Error Err = extractOffloadBundle(Contents, SectionOffset,213 Obj.getFileName(), Bundles))214 return Err;215 }216 }217 return Error::success();218}219 220Error object::extractCodeObject(const ObjectFile &Source, int64_t Offset,221 int64_t Size, StringRef OutputFileName) {222 Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =223 FileOutputBuffer::create(OutputFileName, Size);224 225 if (!BufferOrErr)226 return BufferOrErr.takeError();227 228 Expected<MemoryBufferRef> InputBuffOrErr = Source.getMemoryBufferRef();229 if (Error Err = InputBuffOrErr.takeError())230 return Err;231 232 std::unique_ptr<FileOutputBuffer> Buf = std::move(*BufferOrErr);233 std::copy(InputBuffOrErr->getBufferStart() + Offset,234 InputBuffOrErr->getBufferStart() + Offset + Size,235 Buf->getBufferStart());236 if (Error E = Buf->commit())237 return E;238 239 return Error::success();240}241 242Error object::extractCodeObject(const MemoryBufferRef Buffer, int64_t Offset,243 int64_t Size, StringRef OutputFileName) {244 Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =245 FileOutputBuffer::create(OutputFileName, Size);246 if (!BufferOrErr)247 return BufferOrErr.takeError();248 249 std::unique_ptr<FileOutputBuffer> Buf = std::move(*BufferOrErr);250 std::copy(Buffer.getBufferStart() + Offset,251 Buffer.getBufferStart() + Offset + Size, Buf->getBufferStart());252 253 return Buf->commit();254}255 256// given a file name, offset, and size, extract data into a code object file,257// into file "<SourceFile>-offset<Offset>-size<Size>.co".258Error object::extractOffloadBundleByURI(StringRef URIstr) {259 // create a URI object260 Expected<std::unique_ptr<OffloadBundleURI>> UriOrErr(261 OffloadBundleURI::createOffloadBundleURI(URIstr, FILE_URI));262 if (!UriOrErr)263 return UriOrErr.takeError();264 265 OffloadBundleURI &Uri = **UriOrErr;266 std::string OutputFile = Uri.FileName.str();267 OutputFile +=268 "-offset" + itostr(Uri.Offset) + "-size" + itostr(Uri.Size) + ".co";269 270 // Create an ObjectFile object from uri.file_uri.271 auto ObjOrErr = ObjectFile::createObjectFile(Uri.FileName);272 if (!ObjOrErr)273 return ObjOrErr.takeError();274 275 auto Obj = ObjOrErr->getBinary();276 if (Error Err =277 object::extractCodeObject(*Obj, Uri.Offset, Uri.Size, OutputFile))278 return Err;279 280 return Error::success();281}282 283// Utility function to format numbers with commas.284static std::string formatWithCommas(unsigned long long Value) {285 std::string Num = std::to_string(Value);286 int InsertPosition = Num.length() - 3;287 while (InsertPosition > 0) {288 Num.insert(InsertPosition, ",");289 InsertPosition -= 3;290 }291 return Num;292}293 294Expected<std::unique_ptr<MemoryBuffer>>295CompressedOffloadBundle::compress(compression::Params P,296 const MemoryBuffer &Input, uint16_t Version,297 raw_ostream *VerboseStream) {298 if (!compression::zstd::isAvailable() && !compression::zlib::isAvailable())299 return createStringError("compression not supported.");300 Timer HashTimer("Hash Calculation Timer", "Hash calculation time",301 OffloadBundlerTimerGroup);302 if (VerboseStream)303 HashTimer.startTimer();304 MD5 Hash;305 MD5::MD5Result Result;306 Hash.update(Input.getBuffer());307 Hash.final(Result);308 uint64_t TruncatedHash = Result.low();309 if (VerboseStream)310 HashTimer.stopTimer();311 312 SmallVector<uint8_t, 0> CompressedBuffer;313 auto BufferUint8 = ArrayRef<uint8_t>(314 reinterpret_cast<const uint8_t *>(Input.getBuffer().data()),315 Input.getBuffer().size());316 Timer CompressTimer("Compression Timer", "Compression time",317 OffloadBundlerTimerGroup);318 if (VerboseStream)319 CompressTimer.startTimer();320 compression::compress(P, BufferUint8, CompressedBuffer);321 if (VerboseStream)322 CompressTimer.stopTimer();323 324 uint16_t CompressionMethod = static_cast<uint16_t>(P.format);325 326 // Store sizes in 64-bit variables first.327 uint64_t UncompressedSize64 = Input.getBuffer().size();328 uint64_t TotalFileSize64;329 330 // Calculate total file size based on version.331 if (Version == 2) {332 // For V2, ensure the sizes don't exceed 32-bit limit.333 if (UncompressedSize64 > std::numeric_limits<uint32_t>::max())334 return createStringError("uncompressed size (%llu) exceeds version 2 "335 "unsigned 32-bit integer limit",336 UncompressedSize64);337 TotalFileSize64 = MagicNumber.size() + sizeof(uint32_t) + sizeof(Version) +338 sizeof(CompressionMethod) + sizeof(uint32_t) +339 sizeof(TruncatedHash) + CompressedBuffer.size();340 if (TotalFileSize64 > std::numeric_limits<uint32_t>::max())341 return createStringError("total file size (%llu) exceeds version 2 "342 "unsigned 32-bit integer limit",343 TotalFileSize64);344 345 } else { // Version 3.346 TotalFileSize64 = MagicNumber.size() + sizeof(uint64_t) + sizeof(Version) +347 sizeof(CompressionMethod) + sizeof(uint64_t) +348 sizeof(TruncatedHash) + CompressedBuffer.size();349 }350 351 SmallVector<char, 0> FinalBuffer;352 raw_svector_ostream OS(FinalBuffer);353 OS << MagicNumber;354 OS.write(reinterpret_cast<const char *>(&Version), sizeof(Version));355 OS.write(reinterpret_cast<const char *>(&CompressionMethod),356 sizeof(CompressionMethod));357 358 // Write size fields according to version.359 if (Version == 2) {360 uint32_t TotalFileSize32 = static_cast<uint32_t>(TotalFileSize64);361 uint32_t UncompressedSize32 = static_cast<uint32_t>(UncompressedSize64);362 OS.write(reinterpret_cast<const char *>(&TotalFileSize32),363 sizeof(TotalFileSize32));364 OS.write(reinterpret_cast<const char *>(&UncompressedSize32),365 sizeof(UncompressedSize32));366 } else { // Version 3.367 OS.write(reinterpret_cast<const char *>(&TotalFileSize64),368 sizeof(TotalFileSize64));369 OS.write(reinterpret_cast<const char *>(&UncompressedSize64),370 sizeof(UncompressedSize64));371 }372 373 OS.write(reinterpret_cast<const char *>(&TruncatedHash),374 sizeof(TruncatedHash));375 OS.write(reinterpret_cast<const char *>(CompressedBuffer.data()),376 CompressedBuffer.size());377 378 if (VerboseStream) {379 auto MethodUsed = P.format == compression::Format::Zstd ? "zstd" : "zlib";380 double CompressionRate =381 static_cast<double>(UncompressedSize64) / CompressedBuffer.size();382 double CompressionTimeSeconds = CompressTimer.getTotalTime().getWallTime();383 double CompressionSpeedMBs =384 (UncompressedSize64 / (1024.0 * 1024.0)) / CompressionTimeSeconds;385 *VerboseStream << "Compressed bundle format version: " << Version << "\n"386 << "Total file size (including headers): "387 << formatWithCommas(TotalFileSize64) << " bytes\n"388 << "Compression method used: " << MethodUsed << "\n"389 << "Compression level: " << P.level << "\n"390 << "Binary size before compression: "391 << formatWithCommas(UncompressedSize64) << " bytes\n"392 << "Binary size after compression: "393 << formatWithCommas(CompressedBuffer.size()) << " bytes\n"394 << "Compression rate: " << format("%.2lf", CompressionRate)395 << "\n"396 << "Compression ratio: "397 << format("%.2lf%%", 100.0 / CompressionRate) << "\n"398 << "Compression speed: "399 << format("%.2lf MB/s", CompressionSpeedMBs) << "\n"400 << "Truncated MD5 hash: " << format_hex(TruncatedHash, 16)401 << "\n";402 }403 404 return MemoryBuffer::getMemBufferCopy(405 StringRef(FinalBuffer.data(), FinalBuffer.size()));406}407 408// Use packed structs to avoid padding, such that the structs map the serialized409// format.410LLVM_PACKED_START411union RawCompressedBundleHeader {412 struct CommonFields {413 uint32_t Magic;414 uint16_t Version;415 uint16_t Method;416 };417 418 struct V1Header {419 CommonFields Common;420 uint32_t UncompressedFileSize;421 uint64_t Hash;422 };423 424 struct V2Header {425 CommonFields Common;426 uint32_t FileSize;427 uint32_t UncompressedFileSize;428 uint64_t Hash;429 };430 431 struct V3Header {432 CommonFields Common;433 uint64_t FileSize;434 uint64_t UncompressedFileSize;435 uint64_t Hash;436 };437 438 CommonFields Common;439 V1Header V1;440 V2Header V2;441 V3Header V3;442};443LLVM_PACKED_END444 445// Helper method to get header size based on version.446static size_t getHeaderSize(uint16_t Version) {447 switch (Version) {448 case 1:449 return sizeof(RawCompressedBundleHeader::V1Header);450 case 2:451 return sizeof(RawCompressedBundleHeader::V2Header);452 case 3:453 return sizeof(RawCompressedBundleHeader::V3Header);454 default:455 llvm_unreachable("Unsupported version");456 }457}458 459Expected<CompressedOffloadBundle::CompressedBundleHeader>460CompressedOffloadBundle::CompressedBundleHeader::tryParse(StringRef Blob) {461 assert(Blob.size() >= sizeof(RawCompressedBundleHeader::CommonFields));462 assert(identify_magic(Blob) == file_magic::offload_bundle_compressed);463 464 RawCompressedBundleHeader Header;465 std::memcpy(&Header, Blob.data(), std::min(Blob.size(), sizeof(Header)));466 467 CompressedBundleHeader Normalized;468 Normalized.Version = Header.Common.Version;469 470 size_t RequiredSize = getHeaderSize(Normalized.Version);471 472 if (Blob.size() < RequiredSize)473 return createStringError("compressed bundle header size too small");474 475 switch (Normalized.Version) {476 case 1:477 Normalized.UncompressedFileSize = Header.V1.UncompressedFileSize;478 Normalized.Hash = Header.V1.Hash;479 break;480 case 2:481 Normalized.FileSize = Header.V2.FileSize;482 Normalized.UncompressedFileSize = Header.V2.UncompressedFileSize;483 Normalized.Hash = Header.V2.Hash;484 break;485 case 3:486 Normalized.FileSize = Header.V3.FileSize;487 Normalized.UncompressedFileSize = Header.V3.UncompressedFileSize;488 Normalized.Hash = Header.V3.Hash;489 break;490 default:491 return createStringError("unknown compressed bundle version");492 }493 494 // Determine compression format.495 switch (Header.Common.Method) {496 case static_cast<uint16_t>(compression::Format::Zlib):497 case static_cast<uint16_t>(compression::Format::Zstd):498 Normalized.CompressionFormat =499 static_cast<compression::Format>(Header.Common.Method);500 break;501 default:502 return createStringError("unknown compressing method");503 }504 505 return Normalized;506}507 508Expected<std::unique_ptr<MemoryBuffer>>509CompressedOffloadBundle::decompress(const MemoryBuffer &Input,510 raw_ostream *VerboseStream) {511 StringRef Blob = Input.getBuffer();512 513 // Check minimum header size (using V1 as it's the smallest).514 if (Blob.size() < sizeof(RawCompressedBundleHeader::CommonFields))515 return MemoryBuffer::getMemBufferCopy(Blob);516 517 if (identify_magic(Blob) != file_magic::offload_bundle_compressed) {518 if (VerboseStream)519 *VerboseStream << "Uncompressed bundle\n";520 return MemoryBuffer::getMemBufferCopy(Blob);521 }522 523 Expected<CompressedBundleHeader> HeaderOrErr =524 CompressedBundleHeader::tryParse(Blob);525 if (!HeaderOrErr)526 return HeaderOrErr.takeError();527 528 const CompressedBundleHeader &Normalized = *HeaderOrErr;529 unsigned ThisVersion = Normalized.Version;530 size_t HeaderSize = getHeaderSize(ThisVersion);531 532 compression::Format CompressionFormat = Normalized.CompressionFormat;533 534 size_t TotalFileSize = Normalized.FileSize.value_or(0);535 size_t UncompressedSize = Normalized.UncompressedFileSize;536 auto StoredHash = Normalized.Hash;537 538 Timer DecompressTimer("Decompression Timer", "Decompression time",539 OffloadBundlerTimerGroup);540 if (VerboseStream)541 DecompressTimer.startTimer();542 543 SmallVector<uint8_t, 0> DecompressedData;544 StringRef CompressedData =545 Blob.substr(HeaderSize, TotalFileSize - HeaderSize);546 547 if (Error DecompressionError = compression::decompress(548 CompressionFormat, arrayRefFromStringRef(CompressedData),549 DecompressedData, UncompressedSize))550 return createStringError("could not decompress embedded file contents: " +551 toString(std::move(DecompressionError)));552 553 if (VerboseStream) {554 DecompressTimer.stopTimer();555 556 double DecompressionTimeSeconds =557 DecompressTimer.getTotalTime().getWallTime();558 559 // Recalculate MD5 hash for integrity check.560 Timer HashRecalcTimer("Hash Recalculation Timer", "Hash recalculation time",561 OffloadBundlerTimerGroup);562 HashRecalcTimer.startTimer();563 MD5 Hash;564 MD5::MD5Result Result;565 Hash.update(ArrayRef<uint8_t>(DecompressedData));566 Hash.final(Result);567 uint64_t RecalculatedHash = Result.low();568 HashRecalcTimer.stopTimer();569 bool HashMatch = (StoredHash == RecalculatedHash);570 571 double CompressionRate =572 static_cast<double>(UncompressedSize) / CompressedData.size();573 double DecompressionSpeedMBs =574 (UncompressedSize / (1024.0 * 1024.0)) / DecompressionTimeSeconds;575 576 *VerboseStream << "Compressed bundle format version: " << ThisVersion577 << "\n";578 if (ThisVersion >= 2)579 *VerboseStream << "Total file size (from header): "580 << formatWithCommas(TotalFileSize) << " bytes\n";581 *VerboseStream582 << "Decompression method: "583 << (CompressionFormat == compression::Format::Zlib ? "zlib" : "zstd")584 << "\n"585 << "Size before decompression: "586 << formatWithCommas(CompressedData.size()) << " bytes\n"587 << "Size after decompression: " << formatWithCommas(UncompressedSize)588 << " bytes\n"589 << "Compression rate: " << format("%.2lf", CompressionRate) << "\n"590 << "Compression ratio: " << format("%.2lf%%", 100.0 / CompressionRate)591 << "\n"592 << "Decompression speed: "593 << format("%.2lf MB/s", DecompressionSpeedMBs) << "\n"594 << "Stored hash: " << format_hex(StoredHash, 16) << "\n"595 << "Recalculated hash: " << format_hex(RecalculatedHash, 16) << "\n"596 << "Hashes match: " << (HashMatch ? "Yes" : "No") << "\n";597 }598 599 return MemoryBuffer::getMemBufferCopy(toStringRef(DecompressedData));600}601