clang/lib/Driver/OffloadBundler.cpp

   1 //===- OffloadBundler.cpp - File Bundling and Unbundling ------------------===//
   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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 ///
   9 /// \file
  10 /// This file implements an offload bundling API that bundles different files
  11 /// that relate with the same source code but different targets into a single
  12 /// one. Also the implements the opposite functionality, i.e. unbundle files
  13 /// previous created by this API.
  14 ///
  15 //===----------------------------------------------------------------------===//
  16
  17 #include "clang/Driver/OffloadBundler.h"
  18 #include "clang/Basic/Cuda.h"
  19 #include "clang/Basic/TargetID.h"
  20 #include "clang/Basic/Version.h"
  21 #include "llvm/ADT/ArrayRef.h"
  22 #include "llvm/ADT/SmallString.h"
  23 #include "llvm/ADT/SmallVector.h"
  24 #include "llvm/ADT/StringExtras.h"
  25 #include "llvm/ADT/StringMap.h"
  26 #include "llvm/ADT/StringRef.h"
  27 #include "llvm/BinaryFormat/Magic.h"
  28 #include "llvm/Object/Archive.h"
  29 #include "llvm/Object/ArchiveWriter.h"
  30 #include "llvm/Object/Binary.h"
  31 #include "llvm/Object/ObjectFile.h"
  32 #include "llvm/Support/Casting.h"
  33 #include "llvm/Support/Compression.h"
  34 #include "llvm/Support/Debug.h"
  35 #include "llvm/Support/EndianStream.h"
  36 #include "llvm/Support/Errc.h"
  37 #include "llvm/Support/Error.h"
  38 #include "llvm/Support/ErrorOr.h"
  39 #include "llvm/Support/FileSystem.h"
  40 #include "llvm/Support/MD5.h"
  41 #include "llvm/Support/MemoryBuffer.h"
  42 #include "llvm/Support/Path.h"
  43 #include "llvm/Support/Program.h"
  44 #include "llvm/Support/Signals.h"
  45 #include "llvm/Support/StringSaver.h"
  46 #include "llvm/Support/Timer.h"
  47 #include "llvm/Support/WithColor.h"
  48 #include "llvm/Support/raw_ostream.h"
  49 #include "llvm/TargetParser/Host.h"
  50 #include "llvm/TargetParser/Triple.h"
  51 #include <algorithm>
  52 #include <cassert>
  53 #include <cstddef>
  54 #include <cstdint>
  55 #include <forward_list>
  56 #include <llvm/Support/Process.h>
  57 #include <memory>
  58 #include <set>
  59 #include <string>
  60 #include <system_error>
  61 #include <utility>
  62
  63 using namespace llvm;
  64 using namespace llvm::object;
  65 using namespace clang;
  66
  67 static llvm::TimerGroup
  68     ClangOffloadBundlerTimerGroup("Clang Offload Bundler Timer Group",
  69                                   "Timer group for clang offload bundler");
  70
  71 /// Magic string that marks the existence of offloading data.
  72 #define OFFLOAD_BUNDLER_MAGIC_STR "__CLANG_OFFLOAD_BUNDLE__"
  73
  74 OffloadTargetInfo::OffloadTargetInfo(const StringRef Target,
  75                                      const OffloadBundlerConfig &BC)
  76     : BundlerConfig(BC) {
  77
  78   // TODO: Add error checking from ClangOffloadBundler.cpp
  79   auto TargetFeatures = Target.split(':');
  80   auto TripleOrGPU = TargetFeatures.first.rsplit('-');
  81
  82   if (clang::StringToCudaArch(TripleOrGPU.second) != clang::CudaArch::UNKNOWN) {
  83     auto KindTriple = TripleOrGPU.first.split('-');
  84     this->OffloadKind = KindTriple.first;
  85
  86     // Enforce optional env field to standardize bundles
  87     llvm::Triple t = llvm::Triple(KindTriple.second);
  88     this->Triple = llvm::Triple(t.getArchName(), t.getVendorName(),
  89                                 t.getOSName(), t.getEnvironmentName());
  90
  91     this->TargetID = Target.substr(Target.find(TripleOrGPU.second));
  92   } else {
  93     auto KindTriple = TargetFeatures.first.split('-');
  94     this->OffloadKind = KindTriple.first;
  95
  96     // Enforce optional env field to standardize bundles
  97     llvm::Triple t = llvm::Triple(KindTriple.second);
  98     this->Triple = llvm::Triple(t.getArchName(), t.getVendorName(),
  99                                 t.getOSName(), t.getEnvironmentName());
 100
 101     this->TargetID = "";
 102   }
 103 }
 104
 105 bool OffloadTargetInfo::hasHostKind() const {
 106   return this->OffloadKind == "host";
 107 }
 108
 109 bool OffloadTargetInfo::isOffloadKindValid() const {
 110   return OffloadKind == "host" || OffloadKind == "openmp" ||
 111          OffloadKind == "hip" || OffloadKind == "hipv4";
 112 }
 113
 114 bool OffloadTargetInfo::isOffloadKindCompatible(
 115     const StringRef TargetOffloadKind) const {
 116   if (OffloadKind == TargetOffloadKind)
 117     return true;
 118   if (BundlerConfig.HipOpenmpCompatible) {
 119     bool HIPCompatibleWithOpenMP = OffloadKind.starts_with_insensitive("hip") &&
 120                                    TargetOffloadKind == "openmp";
 121     bool OpenMPCompatibleWithHIP =
 122         OffloadKind == "openmp" &&
 123         TargetOffloadKind.starts_with_insensitive("hip");
 124     return HIPCompatibleWithOpenMP || OpenMPCompatibleWithHIP;
 125   }
 126   return false;
 127 }
 128
 129 bool OffloadTargetInfo::isTripleValid() const {
 130   return !Triple.str().empty() && Triple.getArch() != Triple::UnknownArch;
 131 }
 132
 133 bool OffloadTargetInfo::operator==(const OffloadTargetInfo &Target) const {
 134   return OffloadKind == Target.OffloadKind &&
 135          Triple.isCompatibleWith(Target.Triple) && TargetID == Target.TargetID;
 136 }
 137
 138 std::string OffloadTargetInfo::str() const {
 139   return Twine(OffloadKind + "-" + Triple.str() + "-" + TargetID).str();
 140 }
 141
 142 static StringRef getDeviceFileExtension(StringRef Device,
 143                                         StringRef BundleFileName) {
 144   if (Device.contains("gfx"))
 145     return ".bc";
 146   if (Device.contains("sm_"))
 147     return ".cubin";
 148   return sys::path::extension(BundleFileName);
 149 }
 150
 151 static std::string getDeviceLibraryFileName(StringRef BundleFileName,
 152                                             StringRef Device) {
 153   StringRef LibName = sys::path::stem(BundleFileName);
 154   StringRef Extension = getDeviceFileExtension(Device, BundleFileName);
 155
 156   std::string Result;
 157   Result += LibName;
 158   Result += Extension;
 159   return Result;
 160 }
 161
 162 /// @brief Checks if a code object \p CodeObjectInfo is compatible with a given
 163 /// target \p TargetInfo.
 164 /// @link https://clang.llvm.org/docs/ClangOffloadBundler.html#bundle-entry-id
 165 bool isCodeObjectCompatible(const OffloadTargetInfo &CodeObjectInfo,
 166                             const OffloadTargetInfo &TargetInfo) {
 167
 168   // Compatible in case of exact match.
 169   if (CodeObjectInfo == TargetInfo) {
 170     DEBUG_WITH_TYPE("CodeObjectCompatibility",
 171                     dbgs() << "Compatible: Exact match: \t[CodeObject: "
 172                            << CodeObjectInfo.str()
 173                            << "]\t:\t[Target: " << TargetInfo.str() << "]\n");
 174     return true;
 175   }
 176
 177   // Incompatible if Kinds or Triples mismatch.
 178   if (!CodeObjectInfo.isOffloadKindCompatible(TargetInfo.OffloadKind) ||
 179       !CodeObjectInfo.Triple.isCompatibleWith(TargetInfo.Triple)) {
 180     DEBUG_WITH_TYPE(
 181         "CodeObjectCompatibility",
 182         dbgs() << "Incompatible: Kind/Triple mismatch \t[CodeObject: "
 183                << CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
 184                << "]\n");
 185     return false;
 186   }
 187
 188   // Incompatible if target IDs are incompatible.
 189   if (!clang::isCompatibleTargetID(CodeObjectInfo.TargetID,
 190                                    TargetInfo.TargetID)) {
 191     DEBUG_WITH_TYPE(
 192         "CodeObjectCompatibility",
 193         dbgs() << "Incompatible: target IDs are incompatible \t[CodeObject: "
 194                << CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
 195                << "]\n");
 196     return false;
 197   }
 198
 199   DEBUG_WITH_TYPE(
 200       "CodeObjectCompatibility",
 201       dbgs() << "Compatible: Code Objects are compatible \t[CodeObject: "
 202              << CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
 203              << "]\n");
 204   return true;
 205 }
 206
 207 namespace {
 208 /// Generic file handler interface.
 209 class FileHandler {
 210 public:
 211   struct BundleInfo {
 212     StringRef BundleID;
 213   };
 214
 215   FileHandler() {}
 216
 217   virtual ~FileHandler() {}
 218
 219   /// Update the file handler with information from the header of the bundled
 220   /// file.
 221   virtual Error ReadHeader(MemoryBuffer &Input) = 0;
 222
 223   /// Read the marker of the next bundled to be read in the file. The bundle
 224   /// name is returned if there is one in the file, or `std::nullopt` if there
 225   /// are no more bundles to be read.
 226   virtual Expected<std::optional<StringRef>>
 227   ReadBundleStart(MemoryBuffer &Input) = 0;
 228
 229   /// Read the marker that closes the current bundle.
 230   virtual Error ReadBundleEnd(MemoryBuffer &Input) = 0;
 231
 232   /// Read the current bundle and write the result into the stream \a OS.
 233   virtual Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) = 0;
 234
 235   /// Write the header of the bundled file to \a OS based on the information
 236   /// gathered from \a Inputs.
 237   virtual Error WriteHeader(raw_ostream &OS,
 238                             ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) = 0;
 239
 240   /// Write the marker that initiates a bundle for the triple \a TargetTriple to
 241   /// \a OS.
 242   virtual Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) = 0;
 243
 244   /// Write the marker that closes a bundle for the triple \a TargetTriple to \a
 245   /// OS.
 246   virtual Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) = 0;
 247
 248   /// Write the bundle from \a Input into \a OS.
 249   virtual Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) = 0;
 250
 251   /// Finalize output file.
 252   virtual Error finalizeOutputFile() { return Error::success(); }
 253
 254   /// List bundle IDs in \a Input.
 255   virtual Error listBundleIDs(MemoryBuffer &Input) {
 256     if (Error Err = ReadHeader(Input))
 257       return Err;
 258     return forEachBundle(Input, [&](const BundleInfo &Info) -> Error {
 259       llvm::outs() << Info.BundleID << '\n';
 260       Error Err = listBundleIDsCallback(Input, Info);
 261       if (Err)
 262         return Err;
 263       return Error::success();
 264     });
 265   }
 266
 267   /// For each bundle in \a Input, do \a Func.
 268   Error forEachBundle(MemoryBuffer &Input,
 269                       std::function<Error(const BundleInfo &)> Func) {
 270     while (true) {
 271       Expected<std::optional<StringRef>> CurTripleOrErr =
 272           ReadBundleStart(Input);
 273       if (!CurTripleOrErr)
 274         return CurTripleOrErr.takeError();
 275
 276       // No more bundles.
 277       if (!*CurTripleOrErr)
 278         break;
 279
 280       StringRef CurTriple = **CurTripleOrErr;
 281       assert(!CurTriple.empty());
 282
 283       BundleInfo Info{CurTriple};
 284       if (Error Err = Func(Info))
 285         return Err;
 286     }
 287     return Error::success();
 288   }
 289
 290 protected:
 291   virtual Error listBundleIDsCallback(MemoryBuffer &Input,
 292                                       const BundleInfo &Info) {
 293     return Error::success();
 294   }
 295 };
 296
 297 /// Handler for binary files. The bundled file will have the following format
 298 /// (all integers are stored in little-endian format):
 299 ///
 300 /// "OFFLOAD_BUNDLER_MAGIC_STR" (ASCII encoding of the string)
 301 ///
 302 /// NumberOfOffloadBundles (8-byte integer)
 303 ///
 304 /// OffsetOfBundle1 (8-byte integer)
 305 /// SizeOfBundle1 (8-byte integer)
 306 /// NumberOfBytesInTripleOfBundle1 (8-byte integer)
 307 /// TripleOfBundle1 (byte length defined before)
 308 ///
 309 /// ...
 310 ///
 311 /// OffsetOfBundleN (8-byte integer)
 312 /// SizeOfBundleN (8-byte integer)
 313 /// NumberOfBytesInTripleOfBundleN (8-byte integer)
 314 /// TripleOfBundleN (byte length defined before)
 315 ///
 316 /// Bundle1
 317 /// ...
 318 /// BundleN
 319
 320 /// Read 8-byte integers from a buffer in little-endian format.
 321 static uint64_t Read8byteIntegerFromBuffer(StringRef Buffer, size_t pos) {
 322   return llvm::support::endian::read64le(Buffer.data() + pos);
 323 }
 324
 325 /// Write 8-byte integers to a buffer in little-endian format.
 326 static void Write8byteIntegerToBuffer(raw_ostream &OS, uint64_t Val) {
 327   llvm::support::endian::write(OS, Val, llvm::endianness::little);
 328 }
 329
 330 class BinaryFileHandler final : public FileHandler {
 331   /// Information about the bundles extracted from the header.
 332   struct BinaryBundleInfo final : public BundleInfo {
 333     /// Size of the bundle.
 334     uint64_t Size = 0u;
 335     /// Offset at which the bundle starts in the bundled file.
 336     uint64_t Offset = 0u;
 337
 338     BinaryBundleInfo() {}
 339     BinaryBundleInfo(uint64_t Size, uint64_t Offset)
 340         : Size(Size), Offset(Offset) {}
 341   };
 342
 343   /// Map between a triple and the corresponding bundle information.
 344   StringMap<BinaryBundleInfo> BundlesInfo;
 345
 346   /// Iterator for the bundle information that is being read.
 347   StringMap<BinaryBundleInfo>::iterator CurBundleInfo;
 348   StringMap<BinaryBundleInfo>::iterator NextBundleInfo;
 349
 350   /// Current bundle target to be written.
 351   std::string CurWriteBundleTarget;
 352
 353   /// Configuration options and arrays for this bundler job
 354   const OffloadBundlerConfig &BundlerConfig;
 355
 356 public:
 357   // TODO: Add error checking from ClangOffloadBundler.cpp
 358   BinaryFileHandler(const OffloadBundlerConfig &BC) : BundlerConfig(BC) {}
 359
 360   ~BinaryFileHandler() final {}
 361
 362   Error ReadHeader(MemoryBuffer &Input) final {
 363     StringRef FC = Input.getBuffer();
 364
 365     // Initialize the current bundle with the end of the container.
 366     CurBundleInfo = BundlesInfo.end();
 367
 368     // Check if buffer is smaller than magic string.
 369     size_t ReadChars = sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1;
 370     if (ReadChars > FC.size())
 371       return Error::success();
 372
 373     // Check if no magic was found.
 374     if (llvm::identify_magic(FC) != llvm::file_magic::offload_bundle)
 375       return Error::success();
 376
 377     // Read number of bundles.
 378     if (ReadChars + 8 > FC.size())
 379       return Error::success();
 380
 381     uint64_t NumberOfBundles = Read8byteIntegerFromBuffer(FC, ReadChars);
 382     ReadChars += 8;
 383
 384     // Read bundle offsets, sizes and triples.
 385     for (uint64_t i = 0; i < NumberOfBundles; ++i) {
 386
 387       // Read offset.
 388       if (ReadChars + 8 > FC.size())
 389         return Error::success();
 390
 391       uint64_t Offset = Read8byteIntegerFromBuffer(FC, ReadChars);
 392       ReadChars += 8;
 393
 394       // Read size.
 395       if (ReadChars + 8 > FC.size())
 396         return Error::success();
 397
 398       uint64_t Size = Read8byteIntegerFromBuffer(FC, ReadChars);
 399       ReadChars += 8;
 400
 401       // Read triple size.
 402       if (ReadChars + 8 > FC.size())
 403         return Error::success();
 404
 405       uint64_t TripleSize = Read8byteIntegerFromBuffer(FC, ReadChars);
 406       ReadChars += 8;
 407
 408       // Read triple.
 409       if (ReadChars + TripleSize > FC.size())
 410         return Error::success();
 411
 412       StringRef Triple(&FC.data()[ReadChars], TripleSize);
 413       ReadChars += TripleSize;
 414
 415       // Check if the offset and size make sense.
 416       if (!Offset || Offset + Size > FC.size())
 417         return Error::success();
 418
 419       assert(!BundlesInfo.contains(Triple) && "Triple is duplicated??");
 420       BundlesInfo[Triple] = BinaryBundleInfo(Size, Offset);
 421     }
 422     // Set the iterator to where we will start to read.
 423     CurBundleInfo = BundlesInfo.end();
 424     NextBundleInfo = BundlesInfo.begin();
 425     return Error::success();
 426   }
 427
 428   Expected<std::optional<StringRef>>
 429   ReadBundleStart(MemoryBuffer &Input) final {
 430     if (NextBundleInfo == BundlesInfo.end())
 431       return std::nullopt;
 432     CurBundleInfo = NextBundleInfo++;
 433     return CurBundleInfo->first();
 434   }
 435
 436   Error ReadBundleEnd(MemoryBuffer &Input) final {
 437     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
 438     return Error::success();
 439   }
 440
 441   Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
 442     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
 443     StringRef FC = Input.getBuffer();
 444     OS.write(FC.data() + CurBundleInfo->second.Offset,
 445              CurBundleInfo->second.Size);
 446     return Error::success();
 447   }
 448
 449   Error WriteHeader(raw_ostream &OS,
 450                     ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
 451
 452     // Compute size of the header.
 453     uint64_t HeaderSize = 0;
 454
 455     HeaderSize += sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1;
 456     HeaderSize += 8; // Number of Bundles
 457
 458     for (auto &T : BundlerConfig.TargetNames) {
 459       HeaderSize += 3 * 8; // Bundle offset, Size of bundle and size of triple.
 460       HeaderSize += T.size(); // The triple.
 461     }
 462
 463     // Write to the buffer the header.
 464     OS << OFFLOAD_BUNDLER_MAGIC_STR;
 465
 466     Write8byteIntegerToBuffer(OS, BundlerConfig.TargetNames.size());
 467
 468     unsigned Idx = 0;
 469     for (auto &T : BundlerConfig.TargetNames) {
 470       MemoryBuffer &MB = *Inputs[Idx++];
 471       HeaderSize = alignTo(HeaderSize, BundlerConfig.BundleAlignment);
 472       // Bundle offset.
 473       Write8byteIntegerToBuffer(OS, HeaderSize);
 474       // Size of the bundle (adds to the next bundle's offset)
 475       Write8byteIntegerToBuffer(OS, MB.getBufferSize());
 476       BundlesInfo[T] = BinaryBundleInfo(MB.getBufferSize(), HeaderSize);
 477       HeaderSize += MB.getBufferSize();
 478       // Size of the triple
 479       Write8byteIntegerToBuffer(OS, T.size());
 480       // Triple
 481       OS << T;
 482     }
 483     return Error::success();
 484   }
 485
 486   Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) final {
 487     CurWriteBundleTarget = TargetTriple.str();
 488     return Error::success();
 489   }
 490
 491   Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) final {
 492     return Error::success();
 493   }
 494
 495   Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) final {
 496     auto BI = BundlesInfo[CurWriteBundleTarget];
 497
 498     // Pad with 0 to reach specified offset.
 499     size_t CurrentPos = OS.tell();
 500     size_t PaddingSize = BI.Offset > CurrentPos ? BI.Offset - CurrentPos : 0;
 501     for (size_t I = 0; I < PaddingSize; ++I)
 502       OS.write('\0');
 503     assert(OS.tell() == BI.Offset);
 504
 505     OS.write(Input.getBufferStart(), Input.getBufferSize());
 506
 507     return Error::success();
 508   }
 509 };
 510
 511 // This class implements a list of temporary files that are removed upon
 512 // object destruction.
 513 class TempFileHandlerRAII {
 514 public:
 515   ~TempFileHandlerRAII() {
 516     for (const auto &File : Files)
 517       sys::fs::remove(File);
 518   }
 519
 520   // Creates temporary file with given contents.
 521   Expected<StringRef> Create(std::optional<ArrayRef<char>> Contents) {
 522     SmallString<128u> File;
 523     if (std::error_code EC =
 524             sys::fs::createTemporaryFile("clang-offload-bundler", "tmp", File))
 525       return createFileError(File, EC);
 526     Files.push_front(File);
 527
 528     if (Contents) {
 529       std::error_code EC;
 530       raw_fd_ostream OS(File, EC);
 531       if (EC)
 532         return createFileError(File, EC);
 533       OS.write(Contents->data(), Contents->size());
 534     }
 535     return Files.front().str();
 536   }
 537
 538 private:
 539   std::forward_list<SmallString<128u>> Files;
 540 };
 541
 542 /// Handler for object files. The bundles are organized by sections with a
 543 /// designated name.
 544 ///
 545 /// To unbundle, we just copy the contents of the designated section.
 546 class ObjectFileHandler final : public FileHandler {
 547
 548   /// The object file we are currently dealing with.
 549   std::unique_ptr<ObjectFile> Obj;
 550
 551   /// Return the input file contents.
 552   StringRef getInputFileContents() const { return Obj->getData(); }
 553
 554   /// Return bundle name (<kind>-<triple>) if the provided section is an offload
 555   /// section.
 556   static Expected<std::optional<StringRef>>
 557   IsOffloadSection(SectionRef CurSection) {
 558     Expected<StringRef> NameOrErr = CurSection.getName();
 559     if (!NameOrErr)
 560       return NameOrErr.takeError();
 561
 562     // If it does not start with the reserved suffix, just skip this section.
 563     if (llvm::identify_magic(*NameOrErr) != llvm::file_magic::offload_bundle)
 564       return std::nullopt;
 565
 566     // Return the triple that is right after the reserved prefix.
 567     return NameOrErr->substr(sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1);
 568   }
 569
 570   /// Total number of inputs.
 571   unsigned NumberOfInputs = 0;
 572
 573   /// Total number of processed inputs, i.e, inputs that were already
 574   /// read from the buffers.
 575   unsigned NumberOfProcessedInputs = 0;
 576
 577   /// Iterator of the current and next section.
 578   section_iterator CurrentSection;
 579   section_iterator NextSection;
 580
 581   /// Configuration options and arrays for this bundler job
 582   const OffloadBundlerConfig &BundlerConfig;
 583
 584 public:
 585   // TODO: Add error checking from ClangOffloadBundler.cpp
 586   ObjectFileHandler(std::unique_ptr<ObjectFile> ObjIn,
 587                     const OffloadBundlerConfig &BC)
 588       : Obj(std::move(ObjIn)), CurrentSection(Obj->section_begin()),
 589         NextSection(Obj->section_begin()), BundlerConfig(BC) {}
 590
 591   ~ObjectFileHandler() final {}
 592
 593   Error ReadHeader(MemoryBuffer &Input) final { return Error::success(); }
 594
 595   Expected<std::optional<StringRef>>
 596   ReadBundleStart(MemoryBuffer &Input) final {
 597     while (NextSection != Obj->section_end()) {
 598       CurrentSection = NextSection;
 599       ++NextSection;
 600
 601       // Check if the current section name starts with the reserved prefix. If
 602       // so, return the triple.
 603       Expected<std::optional<StringRef>> TripleOrErr =
 604           IsOffloadSection(*CurrentSection);
 605       if (!TripleOrErr)
 606         return TripleOrErr.takeError();
 607       if (*TripleOrErr)
 608         return **TripleOrErr;
 609     }
 610     return std::nullopt;
 611   }
 612
 613   Error ReadBundleEnd(MemoryBuffer &Input) final { return Error::success(); }
 614
 615   Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
 616     Expected<StringRef> ContentOrErr = CurrentSection->getContents();
 617     if (!ContentOrErr)
 618       return ContentOrErr.takeError();
 619     StringRef Content = *ContentOrErr;
 620
 621     // Copy fat object contents to the output when extracting host bundle.
 622     if (Content.size() == 1u && Content.front() == 0)
 623       Content = StringRef(Input.getBufferStart(), Input.getBufferSize());
 624
 625     OS.write(Content.data(), Content.size());
 626     return Error::success();
 627   }
 628
 629   Error WriteHeader(raw_ostream &OS,
 630                     ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
 631     assert(BundlerConfig.HostInputIndex != ~0u &&
 632            "Host input index not defined.");
 633
 634     // Record number of inputs.
 635     NumberOfInputs = Inputs.size();
 636     return Error::success();
 637   }
 638
 639   Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) final {
 640     ++NumberOfProcessedInputs;
 641     return Error::success();
 642   }
 643
 644   Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) final {
 645     return Error::success();
 646   }
 647
 648   Error finalizeOutputFile() final {
 649     assert(NumberOfProcessedInputs <= NumberOfInputs &&
 650            "Processing more inputs that actually exist!");
 651     assert(BundlerConfig.HostInputIndex != ~0u &&
 652            "Host input index not defined.");
 653
 654     // If this is not the last output, we don't have to do anything.
 655     if (NumberOfProcessedInputs != NumberOfInputs)
 656       return Error::success();
 657
 658     // We will use llvm-objcopy to add target objects sections to the output
 659     // fat object. These sections should have 'exclude' flag set which tells
 660     // link editor to remove them from linker inputs when linking executable or
 661     // shared library.
 662
 663     assert(BundlerConfig.ObjcopyPath != "" &&
 664            "llvm-objcopy path not specified");
 665
 666     // Temporary files that need to be removed.
 667     TempFileHandlerRAII TempFiles;
 668
 669     // Compose llvm-objcopy command line for add target objects' sections with
 670     // appropriate flags.
 671     BumpPtrAllocator Alloc;
 672     StringSaver SS{Alloc};
 673     SmallVector<StringRef, 8u> ObjcopyArgs{"llvm-objcopy"};
 674
 675     for (unsigned I = 0; I < NumberOfInputs; ++I) {
 676       StringRef InputFile = BundlerConfig.InputFileNames[I];
 677       if (I == BundlerConfig.HostInputIndex) {
 678         // Special handling for the host bundle. We do not need to add a
 679         // standard bundle for the host object since we are going to use fat
 680         // object as a host object. Therefore use dummy contents (one zero byte)
 681         // when creating section for the host bundle.
 682         Expected<StringRef> TempFileOrErr = TempFiles.Create(ArrayRef<char>(0));
 683         if (!TempFileOrErr)
 684           return TempFileOrErr.takeError();
 685         InputFile = *TempFileOrErr;
 686       }
 687
 688       ObjcopyArgs.push_back(
 689           SS.save(Twine("--add-section=") + OFFLOAD_BUNDLER_MAGIC_STR +
 690                   BundlerConfig.TargetNames[I] + "=" + InputFile));
 691       ObjcopyArgs.push_back(
 692           SS.save(Twine("--set-section-flags=") + OFFLOAD_BUNDLER_MAGIC_STR +
 693                   BundlerConfig.TargetNames[I] + "=readonly,exclude"));
 694     }
 695     ObjcopyArgs.push_back("--");
 696     ObjcopyArgs.push_back(
 697         BundlerConfig.InputFileNames[BundlerConfig.HostInputIndex]);
 698     ObjcopyArgs.push_back(BundlerConfig.OutputFileNames.front());
 699
 700     if (Error Err = executeObjcopy(BundlerConfig.ObjcopyPath, ObjcopyArgs))
 701       return Err;
 702
 703     return Error::success();
 704   }
 705
 706   Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) final {
 707     return Error::success();
 708   }
 709
 710 private:
 711   Error executeObjcopy(StringRef Objcopy, ArrayRef<StringRef> Args) {
 712     // If the user asked for the commands to be printed out, we do that
 713     // instead of executing it.
 714     if (BundlerConfig.PrintExternalCommands) {
 715       errs() << "\"" << Objcopy << "\"";
 716       for (StringRef Arg : drop_begin(Args, 1))
 717         errs() << " \"" << Arg << "\"";
 718       errs() << "\n";
 719     } else {
 720       if (sys::ExecuteAndWait(Objcopy, Args))
 721         return createStringError(inconvertibleErrorCode(),
 722                                  "'llvm-objcopy' tool failed");
 723     }
 724     return Error::success();
 725   }
 726 };
 727
 728 /// Handler for text files. The bundled file will have the following format.
 729 ///
 730 /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
 731 /// Bundle 1
 732 /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__END__ triple"
 733 /// ...
 734 /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
 735 /// Bundle N
 736 /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__END__ triple"
 737 class TextFileHandler final : public FileHandler {
 738   /// String that begins a line comment.
 739   StringRef Comment;
 740
 741   /// String that initiates a bundle.
 742   std::string BundleStartString;
 743
 744   /// String that closes a bundle.
 745   std::string BundleEndString;
 746
 747   /// Number of chars read from input.
 748   size_t ReadChars = 0u;
 749
 750 protected:
 751   Error ReadHeader(MemoryBuffer &Input) final { return Error::success(); }
 752
 753   Expected<std::optional<StringRef>>
 754   ReadBundleStart(MemoryBuffer &Input) final {
 755     StringRef FC = Input.getBuffer();
 756
 757     // Find start of the bundle.
 758     ReadChars = FC.find(BundleStartString, ReadChars);
 759     if (ReadChars == FC.npos)
 760       return std::nullopt;
 761
 762     // Get position of the triple.
 763     size_t TripleStart = ReadChars = ReadChars + BundleStartString.size();
 764
 765     // Get position that closes the triple.
 766     size_t TripleEnd = ReadChars = FC.find("\n", ReadChars);
 767     if (TripleEnd == FC.npos)
 768       return std::nullopt;
 769
 770     // Next time we read after the new line.
 771     ++ReadChars;
 772
 773     return StringRef(&FC.data()[TripleStart], TripleEnd - TripleStart);
 774   }
 775
 776   Error ReadBundleEnd(MemoryBuffer &Input) final {
 777     StringRef FC = Input.getBuffer();
 778
 779     // Read up to the next new line.
 780     assert(FC[ReadChars] == '\n' && "The bundle should end with a new line.");
 781
 782     size_t TripleEnd = ReadChars = FC.find("\n", ReadChars + 1);
 783     if (TripleEnd != FC.npos)
 784       // Next time we read after the new line.
 785       ++ReadChars;
 786
 787     return Error::success();
 788   }
 789
 790   Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
 791     StringRef FC = Input.getBuffer();
 792     size_t BundleStart = ReadChars;
 793
 794     // Find end of the bundle.
 795     size_t BundleEnd = ReadChars = FC.find(BundleEndString, ReadChars);
 796
 797     StringRef Bundle(&FC.data()[BundleStart], BundleEnd - BundleStart);
 798     OS << Bundle;
 799
 800     return Error::success();
 801   }
 802
 803   Error WriteHeader(raw_ostream &OS,
 804                     ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
 805     return Error::success();
 806   }
 807
 808   Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) final {
 809     OS << BundleStartString << TargetTriple << "\n";
 810     return Error::success();
 811   }
 812
 813   Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) final {
 814     OS << BundleEndString << TargetTriple << "\n";
 815     return Error::success();
 816   }
 817
 818   Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) final {
 819     OS << Input.getBuffer();
 820     return Error::success();
 821   }
 822
 823 public:
 824   TextFileHandler(StringRef Comment) : Comment(Comment), ReadChars(0) {
 825     BundleStartString =
 826         "\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__START__ ";
 827     BundleEndString =
 828         "\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__END__ ";
 829   }
 830
 831   Error listBundleIDsCallback(MemoryBuffer &Input,
 832                               const BundleInfo &Info) final {
 833     // TODO: To list bundle IDs in a bundled text file we need to go through
 834     // all bundles. The format of bundled text file may need to include a
 835     // header if the performance of listing bundle IDs of bundled text file is
 836     // important.
 837     ReadChars = Input.getBuffer().find(BundleEndString, ReadChars);
 838     if (Error Err = ReadBundleEnd(Input))
 839       return Err;
 840     return Error::success();
 841   }
 842 };
 843 } // namespace
 844
 845 /// Return an appropriate object file handler. We use the specific object
 846 /// handler if we know how to deal with that format, otherwise we use a default
 847 /// binary file handler.
 848 static std::unique_ptr<FileHandler>
 849 CreateObjectFileHandler(MemoryBuffer &FirstInput,
 850                         const OffloadBundlerConfig &BundlerConfig) {
 851   // Check if the input file format is one that we know how to deal with.
 852   Expected<std::unique_ptr<Binary>> BinaryOrErr = createBinary(FirstInput);
 853
 854   // We only support regular object files. If failed to open the input as a
 855   // known binary or this is not an object file use the default binary handler.
 856   if (errorToBool(BinaryOrErr.takeError()) || !isa<ObjectFile>(*BinaryOrErr))
 857     return std::make_unique<BinaryFileHandler>(BundlerConfig);
 858
 859   // Otherwise create an object file handler. The handler will be owned by the
 860   // client of this function.
 861   return std::make_unique<ObjectFileHandler>(
 862       std::unique_ptr<ObjectFile>(cast<ObjectFile>(BinaryOrErr->release())),
 863       BundlerConfig);
 864 }
 865
 866 /// Return an appropriate handler given the input files and options.
 867 static Expected<std::unique_ptr<FileHandler>>
 868 CreateFileHandler(MemoryBuffer &FirstInput,
 869                   const OffloadBundlerConfig &BundlerConfig) {
 870   std::string FilesType = BundlerConfig.FilesType;
 871
 872   if (FilesType == "i")
 873     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
 874   if (FilesType == "ii")
 875     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
 876   if (FilesType == "cui")
 877     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
 878   if (FilesType == "hipi")
 879     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
 880   // TODO: `.d` should be eventually removed once `-M` and its variants are
 881   // handled properly in offload compilation.
 882   if (FilesType == "d")
 883     return std::make_unique<TextFileHandler>(/*Comment=*/"#");
 884   if (FilesType == "ll")
 885     return std::make_unique<TextFileHandler>(/*Comment=*/";");
 886   if (FilesType == "bc")
 887     return std::make_unique<BinaryFileHandler>(BundlerConfig);
 888   if (FilesType == "s")
 889     return std::make_unique<TextFileHandler>(/*Comment=*/"#");
 890   if (FilesType == "o")
 891     return CreateObjectFileHandler(FirstInput, BundlerConfig);
 892   if (FilesType == "a")
 893     return CreateObjectFileHandler(FirstInput, BundlerConfig);
 894   if (FilesType == "gch")
 895     return std::make_unique<BinaryFileHandler>(BundlerConfig);
 896   if (FilesType == "ast")
 897     return std::make_unique<BinaryFileHandler>(BundlerConfig);
 898
 899   return createStringError(errc::invalid_argument,
 900                            "'" + FilesType + "': invalid file type specified");
 901 }
 902
 903 OffloadBundlerConfig::OffloadBundlerConfig() {
 904   auto IgnoreEnvVarOpt =
 905       llvm::sys::Process::GetEnv("OFFLOAD_BUNDLER_IGNORE_ENV_VAR");
 906   if (IgnoreEnvVarOpt.has_value() && IgnoreEnvVarOpt.value() == "1")
 907     return;
 908
 909   auto VerboseEnvVarOpt = llvm::sys::Process::GetEnv("OFFLOAD_BUNDLER_VERBOSE");
 910   if (VerboseEnvVarOpt.has_value())
 911     Verbose = VerboseEnvVarOpt.value() == "1";
 912
 913   auto CompressEnvVarOpt =
 914       llvm::sys::Process::GetEnv("OFFLOAD_BUNDLER_COMPRESS");
 915   if (CompressEnvVarOpt.has_value())
 916     Compress = CompressEnvVarOpt.value() == "1";
 917 }
 918
 919 llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
 920 CompressedOffloadBundle::compress(const llvm::MemoryBuffer &Input,
 921                                   bool Verbose) {
 922   llvm::Timer HashTimer("Hash Calculation Timer", "Hash calculation time",
 923                         ClangOffloadBundlerTimerGroup);
 924   if (Verbose)
 925     HashTimer.startTimer();
 926   llvm::MD5 Hash;
 927   llvm::MD5::MD5Result Result;
 928   Hash.update(Input.getBuffer());
 929   Hash.final(Result);
 930   uint64_t TruncatedHash = Result.low();
 931   if (Verbose)
 932     HashTimer.stopTimer();
 933
 934   SmallVector<uint8_t, 0> CompressedBuffer;
 935   auto BufferUint8 = llvm::ArrayRef<uint8_t>(
 936       reinterpret_cast<const uint8_t *>(Input.getBuffer().data()),
 937       Input.getBuffer().size());
 938
 939   llvm::compression::Format CompressionFormat;
 940
 941   if (llvm::compression::zstd::isAvailable())
 942     CompressionFormat = llvm::compression::Format::Zstd;
 943   else if (llvm::compression::zlib::isAvailable())
 944     CompressionFormat = llvm::compression::Format::Zlib;
 945   else
 946     return createStringError(llvm::inconvertibleErrorCode(),
 947                              "Compression not supported");
 948
 949   llvm::Timer CompressTimer("Compression Timer", "Compression time",
 950                             ClangOffloadBundlerTimerGroup);
 951   if (Verbose)
 952     CompressTimer.startTimer();
 953   llvm::compression::compress(CompressionFormat, BufferUint8, CompressedBuffer);
 954   if (Verbose)
 955     CompressTimer.stopTimer();
 956
 957   uint16_t CompressionMethod = static_cast<uint16_t>(CompressionFormat);
 958   uint32_t UncompressedSize = Input.getBuffer().size();
 959
 960   SmallVector<char, 0> FinalBuffer;
 961   llvm::raw_svector_ostream OS(FinalBuffer);
 962   OS << MagicNumber;
 963   OS.write(reinterpret_cast<const char *>(&Version), sizeof(Version));
 964   OS.write(reinterpret_cast<const char *>(&CompressionMethod),
 965            sizeof(CompressionMethod));
 966   OS.write(reinterpret_cast<const char *>(&UncompressedSize),
 967            sizeof(UncompressedSize));
 968   OS.write(reinterpret_cast<const char *>(&TruncatedHash),
 969            sizeof(TruncatedHash));
 970   OS.write(reinterpret_cast<const char *>(CompressedBuffer.data()),
 971            CompressedBuffer.size());
 972
 973   if (Verbose) {
 974     auto MethodUsed =
 975         CompressionFormat == llvm::compression::Format::Zstd ? "zstd" : "zlib";
 976     llvm::errs() << "Compressed bundle format version: " << Version << "\n"
 977                  << "Compression method used: " << MethodUsed << "\n"
 978                  << "Binary size before compression: " << UncompressedSize
 979                  << " bytes\n"
 980                  << "Binary size after compression: " << CompressedBuffer.size()
 981                  << " bytes\n"
 982                  << "Truncated MD5 hash: "
 983                  << llvm::format_hex(TruncatedHash, 16) << "\n";
 984   }
 985
 986   return llvm::MemoryBuffer::getMemBufferCopy(
 987       llvm::StringRef(FinalBuffer.data(), FinalBuffer.size()));
 988 }
 989
 990 llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
 991 CompressedOffloadBundle::decompress(const llvm::MemoryBuffer &Input,
 992                                     bool Verbose) {
 993
 994   StringRef Blob = Input.getBuffer();
 995
 996   if (Blob.size() < HeaderSize) {
 997     return llvm::MemoryBuffer::getMemBufferCopy(Blob);
 998   }
 999   if (llvm::identify_magic(Blob) !=
1000       llvm::file_magic::offload_bundle_compressed) {
1001     if (Verbose)
1002       llvm::errs() << "Uncompressed bundle.\n";
1003     return llvm::MemoryBuffer::getMemBufferCopy(Blob);
1004   }
1005
1006   uint16_t ThisVersion;
1007   uint16_t CompressionMethod;
1008   uint32_t UncompressedSize;
1009   uint64_t StoredHash;
1010   memcpy(&ThisVersion, Input.getBuffer().data() + MagicNumber.size(),
1011          sizeof(uint16_t));
1012   memcpy(&CompressionMethod, Blob.data() + MagicSize + VersionFieldSize,
1013          sizeof(uint16_t));
1014   memcpy(&UncompressedSize,
1015          Blob.data() + MagicSize + VersionFieldSize + MethodFieldSize,
1016          sizeof(uint32_t));
1017   memcpy(&StoredHash,
1018          Blob.data() + MagicSize + VersionFieldSize + MethodFieldSize +
1019              SizeFieldSize,
1020          sizeof(uint64_t));
1021
1022   llvm::compression::Format CompressionFormat;
1023   if (CompressionMethod ==
1024       static_cast<uint16_t>(llvm::compression::Format::Zlib))
1025     CompressionFormat = llvm::compression::Format::Zlib;
1026   else if (CompressionMethod ==
1027            static_cast<uint16_t>(llvm::compression::Format::Zstd))
1028     CompressionFormat = llvm::compression::Format::Zstd;
1029   else
1030     return createStringError(inconvertibleErrorCode(),
1031                              "Unknown compressing method");
1032
1033   llvm::Timer DecompressTimer("Decompression Timer", "Decompression time",
1034                               ClangOffloadBundlerTimerGroup);
1035   if (Verbose)
1036     DecompressTimer.startTimer();
1037
1038   SmallVector<uint8_t, 0> DecompressedData;
1039   StringRef CompressedData = Blob.substr(HeaderSize);
1040   if (llvm::Error DecompressionError = llvm::compression::decompress(
1041           CompressionFormat, llvm::arrayRefFromStringRef(CompressedData),
1042           DecompressedData, UncompressedSize))
1043     return createStringError(inconvertibleErrorCode(),
1044                              "Could not decompress embedded file contents: " +
1045                                  llvm::toString(std::move(DecompressionError)));
1046
1047   if (Verbose) {
1048     DecompressTimer.stopTimer();
1049
1050     // Recalculate MD5 hash
1051     llvm::Timer HashRecalcTimer("Hash Recalculation Timer",
1052                                 "Hash recalculation time",
1053                                 ClangOffloadBundlerTimerGroup);
1054     HashRecalcTimer.startTimer();
1055     llvm::MD5 Hash;
1056     llvm::MD5::MD5Result Result;
1057     Hash.update(llvm::ArrayRef<uint8_t>(DecompressedData.data(),
1058                                         DecompressedData.size()));
1059     Hash.final(Result);
1060     uint64_t RecalculatedHash = Result.low();
1061     HashRecalcTimer.stopTimer();
1062     bool HashMatch = (StoredHash == RecalculatedHash);
1063
1064     llvm::errs() << "Compressed bundle format version: " << ThisVersion << "\n"
1065                  << "Decompression method: "
1066                  << (CompressionFormat == llvm::compression::Format::Zlib
1067                          ? "zlib"
1068                          : "zstd")
1069                  << "\n"
1070                  << "Size before decompression: " << CompressedData.size()
1071                  << " bytes\n"
1072                  << "Size after decompression: " << UncompressedSize
1073                  << " bytes\n"
1074                  << "Stored hash: " << llvm::format_hex(StoredHash, 16) << "\n"
1075                  << "Recalculated hash: "
1076                  << llvm::format_hex(RecalculatedHash, 16) << "\n"
1077                  << "Hashes match: " << (HashMatch ? "Yes" : "No") << "\n";
1078   }
1079
1080   return llvm::MemoryBuffer::getMemBufferCopy(
1081       llvm::toStringRef(DecompressedData));
1082 }
1083
1084 // List bundle IDs. Return true if an error was found.
1085 Error OffloadBundler::ListBundleIDsInFile(
1086     StringRef InputFileName, const OffloadBundlerConfig &BundlerConfig) {
1087   // Open Input file.
1088   ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
1089       MemoryBuffer::getFileOrSTDIN(InputFileName);
1090   if (std::error_code EC = CodeOrErr.getError())
1091     return createFileError(InputFileName, EC);
1092
1093   // Decompress the input if necessary.
1094   Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =
1095       CompressedOffloadBundle::decompress(**CodeOrErr, BundlerConfig.Verbose);
1096   if (!DecompressedBufferOrErr)
1097     return createStringError(
1098         inconvertibleErrorCode(),
1099         "Failed to decompress input: " +
1100             llvm::toString(DecompressedBufferOrErr.takeError()));
1101
1102   MemoryBuffer &DecompressedInput = **DecompressedBufferOrErr;
1103
1104   // Select the right files handler.
1105   Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1106       CreateFileHandler(DecompressedInput, BundlerConfig);
1107   if (!FileHandlerOrErr)
1108     return FileHandlerOrErr.takeError();
1109
1110   std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
1111   assert(FH);
1112   return FH->listBundleIDs(DecompressedInput);
1113 }
1114
1115 /// Bundle the files. Return true if an error was found.
1116 Error OffloadBundler::BundleFiles() {
1117   std::error_code EC;
1118
1119   // Create a buffer to hold the content before compressing.
1120   SmallVector<char, 0> Buffer;
1121   llvm::raw_svector_ostream BufferStream(Buffer);
1122
1123   // Open input files.
1124   SmallVector<std::unique_ptr<MemoryBuffer>, 8u> InputBuffers;
1125   InputBuffers.reserve(BundlerConfig.InputFileNames.size());
1126   for (auto &I : BundlerConfig.InputFileNames) {
1127     ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
1128         MemoryBuffer::getFileOrSTDIN(I);
1129     if (std::error_code EC = CodeOrErr.getError())
1130       return createFileError(I, EC);
1131     InputBuffers.emplace_back(std::move(*CodeOrErr));
1132   }
1133
1134   // Get the file handler. We use the host buffer as reference.
1135   assert((BundlerConfig.HostInputIndex != ~0u || BundlerConfig.AllowNoHost) &&
1136          "Host input index undefined??");
1137   Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr = CreateFileHandler(
1138       *InputBuffers[BundlerConfig.AllowNoHost ? 0
1139                                               : BundlerConfig.HostInputIndex],
1140       BundlerConfig);
1141   if (!FileHandlerOrErr)
1142     return FileHandlerOrErr.takeError();
1143
1144   std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
1145   assert(FH);
1146
1147   // Write header.
1148   if (Error Err = FH->WriteHeader(BufferStream, InputBuffers))
1149     return Err;
1150
1151   // Write all bundles along with the start/end markers. If an error was found
1152   // writing the end of the bundle component, abort the bundle writing.
1153   auto Input = InputBuffers.begin();
1154   for (auto &Triple : BundlerConfig.TargetNames) {
1155     if (Error Err = FH->WriteBundleStart(BufferStream, Triple))
1156       return Err;
1157     if (Error Err = FH->WriteBundle(BufferStream, **Input))
1158       return Err;
1159     if (Error Err = FH->WriteBundleEnd(BufferStream, Triple))
1160       return Err;
1161     ++Input;
1162   }
1163
1164   raw_fd_ostream OutputFile(BundlerConfig.OutputFileNames.front(), EC,
1165                             sys::fs::OF_None);
1166   if (EC)
1167     return createFileError(BundlerConfig.OutputFileNames.front(), EC);
1168
1169   SmallVector<char, 0> CompressedBuffer;
1170   if (BundlerConfig.Compress) {
1171     std::unique_ptr<llvm::MemoryBuffer> BufferMemory =
1172         llvm::MemoryBuffer::getMemBufferCopy(
1173             llvm::StringRef(Buffer.data(), Buffer.size()));
1174     auto CompressionResult =
1175         CompressedOffloadBundle::compress(*BufferMemory, BundlerConfig.Verbose);
1176     if (auto Error = CompressionResult.takeError())
1177       return Error;
1178
1179     auto CompressedMemBuffer = std::move(CompressionResult.get());
1180     CompressedBuffer.assign(CompressedMemBuffer->getBufferStart(),
1181                             CompressedMemBuffer->getBufferEnd());
1182   } else
1183     CompressedBuffer = Buffer;
1184
1185   OutputFile.write(CompressedBuffer.data(), CompressedBuffer.size());
1186
1187   return FH->finalizeOutputFile();
1188 }
1189
1190 // Unbundle the files. Return true if an error was found.
1191 Error OffloadBundler::UnbundleFiles() {
1192   // Open Input file.
1193   ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
1194       MemoryBuffer::getFileOrSTDIN(BundlerConfig.InputFileNames.front());
1195   if (std::error_code EC = CodeOrErr.getError())
1196     return createFileError(BundlerConfig.InputFileNames.front(), EC);
1197
1198   // Decompress the input if necessary.
1199   Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =
1200       CompressedOffloadBundle::decompress(**CodeOrErr, BundlerConfig.Verbose);
1201   if (!DecompressedBufferOrErr)
1202     return createStringError(
1203         inconvertibleErrorCode(),
1204         "Failed to decompress input: " +
1205             llvm::toString(DecompressedBufferOrErr.takeError()));
1206
1207   MemoryBuffer &Input = **DecompressedBufferOrErr;
1208
1209   // Select the right files handler.
1210   Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1211       CreateFileHandler(Input, BundlerConfig);
1212   if (!FileHandlerOrErr)
1213     return FileHandlerOrErr.takeError();
1214
1215   std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
1216   assert(FH);
1217
1218   // Read the header of the bundled file.
1219   if (Error Err = FH->ReadHeader(Input))
1220     return Err;
1221
1222   // Create a work list that consist of the map triple/output file.
1223   StringMap<StringRef> Worklist;
1224   auto Output = BundlerConfig.OutputFileNames.begin();
1225   for (auto &Triple : BundlerConfig.TargetNames) {
1226     Worklist[Triple] = *Output;
1227     ++Output;
1228   }
1229
1230   // Read all the bundles that are in the work list. If we find no bundles we
1231   // assume the file is meant for the host target.
1232   bool FoundHostBundle = false;
1233   while (!Worklist.empty()) {
1234     Expected<std::optional<StringRef>> CurTripleOrErr =
1235         FH->ReadBundleStart(Input);
1236     if (!CurTripleOrErr)
1237       return CurTripleOrErr.takeError();
1238
1239     // We don't have more bundles.
1240     if (!*CurTripleOrErr)
1241       break;
1242
1243     StringRef CurTriple = **CurTripleOrErr;
1244     assert(!CurTriple.empty());
1245
1246     auto Output = Worklist.begin();
1247     for (auto E = Worklist.end(); Output != E; Output++) {
1248       if (isCodeObjectCompatible(
1249               OffloadTargetInfo(CurTriple, BundlerConfig),
1250               OffloadTargetInfo((*Output).first(), BundlerConfig))) {
1251         break;
1252       }
1253     }
1254
1255     if (Output == Worklist.end())
1256       continue;
1257     // Check if the output file can be opened and copy the bundle to it.
1258     std::error_code EC;
1259     raw_fd_ostream OutputFile((*Output).second, EC, sys::fs::OF_None);
1260     if (EC)
1261       return createFileError((*Output).second, EC);
1262     if (Error Err = FH->ReadBundle(OutputFile, Input))
1263       return Err;
1264     if (Error Err = FH->ReadBundleEnd(Input))
1265       return Err;
1266     Worklist.erase(Output);
1267
1268     // Record if we found the host bundle.
1269     auto OffloadInfo = OffloadTargetInfo(CurTriple, BundlerConfig);
1270     if (OffloadInfo.hasHostKind())
1271       FoundHostBundle = true;
1272   }
1273
1274   if (!BundlerConfig.AllowMissingBundles && !Worklist.empty()) {
1275     std::string ErrMsg = "Can't find bundles for";
1276     std::set<StringRef> Sorted;
1277     for (auto &E : Worklist)
1278       Sorted.insert(E.first());
1279     unsigned I = 0;
1280     unsigned Last = Sorted.size() - 1;
1281     for (auto &E : Sorted) {
1282       if (I != 0 && Last > 1)
1283         ErrMsg += ",";
1284       ErrMsg += " ";
1285       if (I == Last && I != 0)
1286         ErrMsg += "and ";
1287       ErrMsg += E.str();
1288       ++I;
1289     }
1290     return createStringError(inconvertibleErrorCode(), ErrMsg);
1291   }
1292
1293   // If no bundles were found, assume the input file is the host bundle and
1294   // create empty files for the remaining targets.
1295   if (Worklist.size() == BundlerConfig.TargetNames.size()) {
1296     for (auto &E : Worklist) {
1297       std::error_code EC;
1298       raw_fd_ostream OutputFile(E.second, EC, sys::fs::OF_None);
1299       if (EC)
1300         return createFileError(E.second, EC);
1301
1302       // If this entry has a host kind, copy the input file to the output file.
1303       auto OffloadInfo = OffloadTargetInfo(E.getKey(), BundlerConfig);
1304       if (OffloadInfo.hasHostKind())
1305         OutputFile.write(Input.getBufferStart(), Input.getBufferSize());
1306     }
1307     return Error::success();
1308   }
1309
1310   // If we found elements, we emit an error if none of those were for the host
1311   // in case host bundle name was provided in command line.
1312   if (!(FoundHostBundle || BundlerConfig.HostInputIndex == ~0u ||
1313         BundlerConfig.AllowMissingBundles))
1314     return createStringError(inconvertibleErrorCode(),
1315                              "Can't find bundle for the host target");
1316
1317   // If we still have any elements in the worklist, create empty files for them.
1318   for (auto &E : Worklist) {
1319     std::error_code EC;
1320     raw_fd_ostream OutputFile(E.second, EC, sys::fs::OF_None);
1321     if (EC)
1322       return createFileError(E.second, EC);
1323   }
1324
1325   return Error::success();
1326 }
1327
1328 static Archive::Kind getDefaultArchiveKindForHost() {
1329   return Triple(sys::getDefaultTargetTriple()).isOSDarwin() ? Archive::K_DARWIN
1330                                                             : Archive::K_GNU;
1331 }
1332
1333 /// @brief Computes a list of targets among all given targets which are
1334 /// compatible with this code object
1335 /// @param [in] CodeObjectInfo Code Object
1336 /// @param [out] CompatibleTargets List of all compatible targets among all
1337 /// given targets
1338 /// @return false, if no compatible target is found.
1339 static bool
1340 getCompatibleOffloadTargets(OffloadTargetInfo &CodeObjectInfo,
1341                             SmallVectorImpl<StringRef> &CompatibleTargets,
1342                             const OffloadBundlerConfig &BundlerConfig) {
1343   if (!CompatibleTargets.empty()) {
1344     DEBUG_WITH_TYPE("CodeObjectCompatibility",
1345                     dbgs() << "CompatibleTargets list should be empty\n");
1346     return false;
1347   }
1348   for (auto &Target : BundlerConfig.TargetNames) {
1349     auto TargetInfo = OffloadTargetInfo(Target, BundlerConfig);
1350     if (isCodeObjectCompatible(CodeObjectInfo, TargetInfo))
1351       CompatibleTargets.push_back(Target);
1352   }
1353   return !CompatibleTargets.empty();
1354 }
1355
1356 /// UnbundleArchive takes an archive file (".a") as input containing bundled
1357 /// code object files, and a list of offload targets (not host), and extracts
1358 /// the code objects into a new archive file for each offload target. Each
1359 /// resulting archive file contains all code object files corresponding to that
1360 /// particular offload target. The created archive file does not
1361 /// contain an index of the symbols and code object files are named as
1362 /// <<Parent Bundle Name>-<CodeObject's GPUArch>>, with ':' replaced with '_'.
1363 Error OffloadBundler::UnbundleArchive() {
1364   std::vector<std::unique_ptr<MemoryBuffer>> ArchiveBuffers;
1365
1366   /// Map of target names with list of object files that will form the device
1367   /// specific archive for that target
1368   StringMap<std::vector<NewArchiveMember>> OutputArchivesMap;
1369
1370   // Map of target names and output archive filenames
1371   StringMap<StringRef> TargetOutputFileNameMap;
1372
1373   auto Output = BundlerConfig.OutputFileNames.begin();
1374   for (auto &Target : BundlerConfig.TargetNames) {
1375     TargetOutputFileNameMap[Target] = *Output;
1376     ++Output;
1377   }
1378
1379   StringRef IFName = BundlerConfig.InputFileNames.front();
1380
1381   ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
1382       MemoryBuffer::getFileOrSTDIN(IFName, true, false);
1383   if (std::error_code EC = BufOrErr.getError())
1384     return createFileError(BundlerConfig.InputFileNames.front(), EC);
1385
1386   ArchiveBuffers.push_back(std::move(*BufOrErr));
1387   Expected<std::unique_ptr<llvm::object::Archive>> LibOrErr =
1388       Archive::create(ArchiveBuffers.back()->getMemBufferRef());
1389   if (!LibOrErr)
1390     return LibOrErr.takeError();
1391
1392   auto Archive = std::move(*LibOrErr);
1393
1394   Error ArchiveErr = Error::success();
1395   auto ChildEnd = Archive->child_end();
1396
1397   /// Iterate over all bundled code object files in the input archive.
1398   for (auto ArchiveIter = Archive->child_begin(ArchiveErr);
1399        ArchiveIter != ChildEnd; ++ArchiveIter) {
1400     if (ArchiveErr)
1401       return ArchiveErr;
1402     auto ArchiveChildNameOrErr = (*ArchiveIter).getName();
1403     if (!ArchiveChildNameOrErr)
1404       return ArchiveChildNameOrErr.takeError();
1405
1406     StringRef BundledObjectFile = sys::path::filename(*ArchiveChildNameOrErr);
1407
1408     auto CodeObjectBufferRefOrErr = (*ArchiveIter).getMemoryBufferRef();
1409     if (!CodeObjectBufferRefOrErr)
1410       return CodeObjectBufferRefOrErr.takeError();
1411
1412     auto TempCodeObjectBuffer =
1413         MemoryBuffer::getMemBuffer(*CodeObjectBufferRefOrErr, false);
1414
1415     // Decompress the buffer if necessary.
1416     Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =
1417         CompressedOffloadBundle::decompress(*TempCodeObjectBuffer,
1418                                             BundlerConfig.Verbose);
1419     if (!DecompressedBufferOrErr)
1420       return createStringError(
1421           inconvertibleErrorCode(),
1422           "Failed to decompress code object: " +
1423               llvm::toString(DecompressedBufferOrErr.takeError()));
1424
1425     MemoryBuffer &CodeObjectBuffer = **DecompressedBufferOrErr;
1426
1427     Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1428         CreateFileHandler(CodeObjectBuffer, BundlerConfig);
1429     if (!FileHandlerOrErr)
1430       return FileHandlerOrErr.takeError();
1431
1432     std::unique_ptr<FileHandler> &FileHandler = *FileHandlerOrErr;
1433     assert(FileHandler &&
1434            "FileHandle creation failed for file in the archive!");
1435
1436     if (Error ReadErr = FileHandler->ReadHeader(CodeObjectBuffer))
1437       return ReadErr;
1438
1439     Expected<std::optional<StringRef>> CurBundleIDOrErr =
1440         FileHandler->ReadBundleStart(CodeObjectBuffer);
1441     if (!CurBundleIDOrErr)
1442       return CurBundleIDOrErr.takeError();
1443
1444     std::optional<StringRef> OptionalCurBundleID = *CurBundleIDOrErr;
1445     // No device code in this child, skip.
1446     if (!OptionalCurBundleID)
1447       continue;
1448     StringRef CodeObject = *OptionalCurBundleID;
1449
1450     // Process all bundle entries (CodeObjects) found in this child of input
1451     // archive.
1452     while (!CodeObject.empty()) {
1453       SmallVector<StringRef> CompatibleTargets;
1454       auto CodeObjectInfo = OffloadTargetInfo(CodeObject, BundlerConfig);
1455       if (CodeObjectInfo.hasHostKind()) {
1456         // Do nothing, we don't extract host code yet.
1457       } else if (getCompatibleOffloadTargets(CodeObjectInfo, CompatibleTargets,
1458                                              BundlerConfig)) {
1459         std::string BundleData;
1460         raw_string_ostream DataStream(BundleData);
1461         if (Error Err = FileHandler->ReadBundle(DataStream, CodeObjectBuffer))
1462           return Err;
1463
1464         for (auto &CompatibleTarget : CompatibleTargets) {
1465           SmallString<128> BundledObjectFileName;
1466           BundledObjectFileName.assign(BundledObjectFile);
1467           auto OutputBundleName =
1468               Twine(llvm::sys::path::stem(BundledObjectFileName) + "-" +
1469                     CodeObject +
1470                     getDeviceLibraryFileName(BundledObjectFileName,
1471                                              CodeObjectInfo.TargetID))
1472                   .str();
1473           // Replace ':' in optional target feature list with '_' to ensure
1474           // cross-platform validity.
1475           std::replace(OutputBundleName.begin(), OutputBundleName.end(), ':',
1476                        '_');
1477
1478           std::unique_ptr<MemoryBuffer> MemBuf = MemoryBuffer::getMemBufferCopy(
1479               DataStream.str(), OutputBundleName);
1480           ArchiveBuffers.push_back(std::move(MemBuf));
1481           llvm::MemoryBufferRef MemBufRef =
1482               MemoryBufferRef(*(ArchiveBuffers.back()));
1483
1484           // For inserting <CompatibleTarget, list<CodeObject>> entry in
1485           // OutputArchivesMap.
1486           if (!OutputArchivesMap.contains(CompatibleTarget)) {
1487
1488             std::vector<NewArchiveMember> ArchiveMembers;
1489             ArchiveMembers.push_back(NewArchiveMember(MemBufRef));
1490             OutputArchivesMap.insert_or_assign(CompatibleTarget,
1491                                                std::move(ArchiveMembers));
1492           } else {
1493             OutputArchivesMap[CompatibleTarget].push_back(
1494                 NewArchiveMember(MemBufRef));
1495           }
1496         }
1497       }
1498
1499       if (Error Err = FileHandler->ReadBundleEnd(CodeObjectBuffer))
1500         return Err;
1501
1502       Expected<std::optional<StringRef>> NextTripleOrErr =
1503           FileHandler->ReadBundleStart(CodeObjectBuffer);
1504       if (!NextTripleOrErr)
1505         return NextTripleOrErr.takeError();
1506
1507       CodeObject = ((*NextTripleOrErr).has_value()) ? **NextTripleOrErr : "";
1508     } // End of processing of all bundle entries of this child of input archive.
1509   }   // End of while over children of input archive.
1510
1511   assert(!ArchiveErr && "Error occurred while reading archive!");
1512
1513   /// Write out an archive for each target
1514   for (auto &Target : BundlerConfig.TargetNames) {
1515     StringRef FileName = TargetOutputFileNameMap[Target];
1516     StringMapIterator<std::vector<llvm::NewArchiveMember>> CurArchiveMembers =
1517         OutputArchivesMap.find(Target);
1518     if (CurArchiveMembers != OutputArchivesMap.end()) {
1519       if (Error WriteErr = writeArchive(FileName, CurArchiveMembers->getValue(),
1520                                         SymtabWritingMode::NormalSymtab,
1521                                         getDefaultArchiveKindForHost(), true,
1522                                         false, nullptr))
1523         return WriteErr;
1524     } else if (!BundlerConfig.AllowMissingBundles) {
1525       std::string ErrMsg =
1526           Twine("no compatible code object found for the target '" + Target +
1527                 "' in heterogeneous archive library: " + IFName)
1528               .str();
1529       return createStringError(inconvertibleErrorCode(), ErrMsg);
1530     } else { // Create an empty archive file if no compatible code object is
1531              // found and "allow-missing-bundles" is enabled. It ensures that
1532              // the linker using output of this step doesn't complain about
1533              // the missing input file.
1534       std::vector<llvm::NewArchiveMember> EmptyArchive;
1535       EmptyArchive.clear();
1536       if (Error WriteErr = writeArchive(
1537               FileName, EmptyArchive, SymtabWritingMode::NormalSymtab,
1538               getDefaultArchiveKindForHost(), true, false, nullptr))
1539         return WriteErr;
1540     }
1541   }
1542
1543   return Error::success();
1544 }