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