1 //===-- clang-linker-wrapper/ClangLinkerWrapper.cpp - wrapper over linker-===//
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 // This tool works as a wrapper over a linking job. This tool is used to create
10 // linked device images for offloading. It scans the linker's input for embedded
11 // device offloading data stored in sections `.llvm.offloading.<triple>.<arch>`
12 // and extracts it as a temporary file. The extracted device files will then be
13 // passed to a device linking job to create a final device image.
14 //
15 //===---------------------------------------------------------------------===//
16 
17 #include "OffloadWrapper.h"
18 #include "clang/Basic/Version.h"
19 #include "llvm/BinaryFormat/Magic.h"
20 #include "llvm/Bitcode/BitcodeWriter.h"
21 #include "llvm/CodeGen/CommandFlags.h"
22 #include "llvm/IR/Constants.h"
23 #include "llvm/IR/DiagnosticPrinter.h"
24 #include "llvm/IR/Module.h"
25 #include "llvm/IRReader/IRReader.h"
26 #include "llvm/LTO/LTO.h"
27 #include "llvm/MC/TargetRegistry.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/Object/OffloadBinary.h"
33 #include "llvm/Support/CommandLine.h"
34 #include "llvm/Support/Errc.h"
35 #include "llvm/Support/FileOutputBuffer.h"
36 #include "llvm/Support/FileSystem.h"
37 #include "llvm/Support/Host.h"
38 #include "llvm/Support/InitLLVM.h"
39 #include "llvm/Support/MemoryBuffer.h"
40 #include "llvm/Support/Path.h"
41 #include "llvm/Support/Program.h"
42 #include "llvm/Support/Signals.h"
43 #include "llvm/Support/SourceMgr.h"
44 #include "llvm/Support/StringSaver.h"
45 #include "llvm/Support/TargetSelect.h"
46 #include "llvm/Support/WithColor.h"
47 #include "llvm/Support/raw_ostream.h"
48 #include "llvm/Target/TargetMachine.h"
49 
50 using namespace llvm;
51 using namespace llvm::object;
52 
53 static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
54 
55 enum DebugKind {
56   NoDebugInfo,
57   DirectivesOnly,
58   FullDebugInfo,
59 };
60 
61 // Mark all our options with this category, everything else (except for -help)
62 // will be hidden.
63 static cl::OptionCategory
64     ClangLinkerWrapperCategory("clang-linker-wrapper options");
65 
66 static cl::opt<bool> StripSections(
67     "strip-sections", cl::ZeroOrMore,
68     cl::desc("Strip offloading sections from the host object file."),
69     cl::init(false), cl::cat(ClangLinkerWrapperCategory));
70 
71 static cl::opt<std::string> LinkerUserPath("linker-path", cl::Required,
72                                            cl::desc("Path of linker binary"),
73                                            cl::cat(ClangLinkerWrapperCategory));
74 
75 static cl::opt<std::string>
76     TargetFeatures("target-feature", cl::ZeroOrMore,
77                    cl::desc("Target features for triple"),
78                    cl::cat(ClangLinkerWrapperCategory));
79 
80 static cl::opt<std::string> OptLevel("opt-level", cl::ZeroOrMore,
81                                      cl::desc("Optimization level for LTO"),
82                                      cl::init("O2"),
83                                      cl::cat(ClangLinkerWrapperCategory));
84 
85 static cl::list<std::string>
86     BitcodeLibraries("target-library", cl::ZeroOrMore,
87                      cl::desc("Path for the target bitcode library"),
88                      cl::cat(ClangLinkerWrapperCategory));
89 
90 static cl::opt<bool> EmbedBitcode(
91     "target-embed-bc", cl::ZeroOrMore,
92     cl::desc("Embed linked bitcode instead of an executable device image"),
93     cl::init(false), cl::cat(ClangLinkerWrapperCategory));
94 
95 static cl::opt<bool> DryRun(
96     "dry-run", cl::ZeroOrMore,
97     cl::desc("List the linker commands to be run without executing them"),
98     cl::init(false), cl::cat(ClangLinkerWrapperCategory));
99 
100 static cl::opt<bool>
101     PrintWrappedModule("print-wrapped-module", cl::ZeroOrMore,
102                        cl::desc("Print the wrapped module's IR for testing"),
103                        cl::init(false), cl::cat(ClangLinkerWrapperCategory));
104 
105 static cl::opt<std::string>
106     HostTriple("host-triple", cl::ZeroOrMore,
107                cl::desc("Triple to use for the host compilation"),
108                cl::init(sys::getDefaultTargetTriple()),
109                cl::cat(ClangLinkerWrapperCategory));
110 
111 static cl::list<std::string>
112     PtxasArgs("ptxas-args", cl::ZeroOrMore,
113               cl::desc("Argument to pass to the ptxas invocation"),
114               cl::cat(ClangLinkerWrapperCategory));
115 
116 static cl::opt<bool> Verbose("v", cl::ZeroOrMore,
117                              cl::desc("Verbose output from tools"),
118                              cl::init(false),
119                              cl::cat(ClangLinkerWrapperCategory));
120 
121 static cl::opt<DebugKind> DebugInfo(
122     cl::desc("Choose debugging level:"), cl::init(NoDebugInfo),
123     cl::values(clEnumValN(NoDebugInfo, "g0", "No debug information"),
124                clEnumValN(DirectivesOnly, "gline-directives-only",
125                           "Direction information"),
126                clEnumValN(FullDebugInfo, "g", "Full debugging support")));
127 
128 static cl::opt<bool> SaveTemps("save-temps", cl::ZeroOrMore,
129                                cl::desc("Save intermediary results."),
130                                cl::cat(ClangLinkerWrapperCategory));
131 
132 static cl::opt<std::string> CudaPath("cuda-path", cl::ZeroOrMore,
133                                      cl::desc("Save intermediary results."),
134                                      cl::cat(ClangLinkerWrapperCategory));
135 
136 // Do not parse linker options.
137 static cl::list<std::string>
138     HostLinkerArgs(cl::Positional,
139                    cl::desc("<options to be passed to linker>..."));
140 
141 /// Path of the current binary.
142 static const char *LinkerExecutable;
143 
144 /// Filename of the executable being created.
145 static StringRef ExecutableName;
146 
147 /// System root if passed in to the linker via. '--sysroot='.
148 static StringRef Sysroot = "";
149 
150 /// Binary path for the CUDA installation.
151 static std::string CudaBinaryPath;
152 
153 /// Temporary files created by the linker wrapper.
154 static SmallVector<std::string, 16> TempFiles;
155 
156 /// Codegen flags for LTO backend.
157 static codegen::RegisterCodeGenFlags CodeGenFlags;
158 
159 /// Magic section string that marks the existence of offloading data. The
160 /// section will contain one or more offloading binaries stored contiguously.
161 #define OFFLOAD_SECTION_MAGIC_STR ".llvm.offloading"
162 
163 /// The magic offset for the first object inside CUDA's fatbinary. This can be
164 /// different but it should work for what is passed here.
165 static constexpr unsigned FatbinaryOffset = 0x50;
166 
167 /// Information for a device offloading file extracted from the host.
168 struct DeviceFile {
169   DeviceFile(StringRef Kind, StringRef TheTriple, StringRef Arch,
170              StringRef Filename, bool IsLibrary = false)
171       : Kind(Kind), TheTriple(TheTriple), Arch(Arch), Filename(Filename),
172         IsLibrary(IsLibrary) {}
173 
174   std::string Kind;
175   std::string TheTriple;
176   std::string Arch;
177   std::string Filename;
178   bool IsLibrary;
179 };
180 
181 namespace llvm {
182 /// Helper that allows DeviceFile to be used as a key in a DenseMap. For now we
183 /// assume device files with matching architectures and triples but different
184 /// offloading kinds should be handlded together, this may not be true in the
185 /// future.
186 template <> struct DenseMapInfo<DeviceFile> {
187   static DeviceFile getEmptyKey() {
188     return {DenseMapInfo<StringRef>::getEmptyKey(),
189             DenseMapInfo<StringRef>::getEmptyKey(),
190             DenseMapInfo<StringRef>::getEmptyKey(),
191             DenseMapInfo<StringRef>::getEmptyKey()};
192   }
193   static DeviceFile getTombstoneKey() {
194     return {DenseMapInfo<StringRef>::getTombstoneKey(),
195             DenseMapInfo<StringRef>::getTombstoneKey(),
196             DenseMapInfo<StringRef>::getTombstoneKey(),
197             DenseMapInfo<StringRef>::getTombstoneKey()};
198   }
199   static unsigned getHashValue(const DeviceFile &I) {
200     return DenseMapInfo<StringRef>::getHashValue(I.TheTriple) ^
201            DenseMapInfo<StringRef>::getHashValue(I.Arch);
202   }
203   static bool isEqual(const DeviceFile &LHS, const DeviceFile &RHS) {
204     return LHS.TheTriple == RHS.TheTriple && LHS.Arch == RHS.Arch;
205   }
206 };
207 } // namespace llvm
208 
209 namespace {
210 
211 Expected<Optional<std::string>>
212 extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
213                   SmallVectorImpl<DeviceFile> &DeviceFiles,
214                   bool IsLibrary = false);
215 
216 void printCommands(ArrayRef<StringRef> CmdArgs) {
217   if (CmdArgs.empty())
218     return;
219 
220   llvm::errs() << " \"" << CmdArgs.front() << "\" ";
221   for (auto IC = std::next(CmdArgs.begin()), IE = CmdArgs.end(); IC != IE; ++IC)
222     llvm::errs() << *IC << (std::next(IC) != IE ? " " : "\n");
223 }
224 
225 std::string getMainExecutable(const char *Name) {
226   void *Ptr = (void *)(intptr_t)&getMainExecutable;
227   auto COWPath = sys::fs::getMainExecutable(Name, Ptr);
228   return sys::path::parent_path(COWPath).str();
229 }
230 
231 /// Extract the device file from the string '<kind>-<triple>-<arch>=<library>'.
232 DeviceFile getBitcodeLibrary(StringRef LibraryStr) {
233   auto DeviceAndPath = StringRef(LibraryStr).split('=');
234   auto StringAndArch = DeviceAndPath.first.rsplit('-');
235   auto KindAndTriple = StringAndArch.first.split('-');
236   return DeviceFile(KindAndTriple.first, KindAndTriple.second,
237                     StringAndArch.second, DeviceAndPath.second);
238 }
239 
240 /// Get a temporary filename suitable for output.
241 Error createOutputFile(const Twine &Prefix, StringRef Extension,
242                        SmallString<128> &NewFilename) {
243   if (!SaveTemps) {
244     if (std::error_code EC =
245             sys::fs::createTemporaryFile(Prefix, Extension, NewFilename))
246       return createFileError(NewFilename, EC);
247     TempFiles.push_back(static_cast<std::string>(NewFilename));
248   } else {
249     const Twine &Filename = Prefix + "." + Extension;
250     Filename.toNullTerminatedStringRef(NewFilename);
251   }
252 
253   return Error::success();
254 }
255 
256 /// Execute the command \p ExecutablePath with the arguments \p Args.
257 Error executeCommands(StringRef ExecutablePath, ArrayRef<StringRef> Args) {
258   if (Verbose || DryRun)
259     printCommands(Args);
260 
261   if (!DryRun)
262     if (sys::ExecuteAndWait(ExecutablePath, Args))
263       return createStringError(inconvertibleErrorCode(),
264                                "'" + sys::path::filename(ExecutablePath) + "'" +
265                                    " failed");
266   return Error::success();
267 }
268 
269 Expected<std::string> findProgram(StringRef Name, ArrayRef<StringRef> Paths) {
270 
271   ErrorOr<std::string> Path = sys::findProgramByName(Name, Paths);
272   if (!Path)
273     Path = sys::findProgramByName(Name);
274   if (!Path && DryRun)
275     return Name.str();
276   if (!Path)
277     return createStringError(Path.getError(),
278                              "Unable to find '" + Name + "' in path");
279   return *Path;
280 }
281 
282 Error runLinker(std::string &LinkerPath, SmallVectorImpl<std::string> &Args) {
283   std::vector<StringRef> LinkerArgs;
284   LinkerArgs.push_back(LinkerPath);
285   for (auto &Arg : Args)
286     LinkerArgs.push_back(Arg);
287 
288   if (Error Err = executeCommands(LinkerPath, LinkerArgs))
289     return Err;
290   return Error::success();
291 }
292 
293 void PrintVersion(raw_ostream &OS) {
294   OS << clang::getClangToolFullVersion("clang-linker-wrapper") << '\n';
295 }
296 
297 void removeFromCompilerUsed(Module &M, GlobalValue &Value) {
298   GlobalVariable *GV = M.getGlobalVariable("llvm.compiler.used");
299   Type *Int8PtrTy = Type::getInt8PtrTy(M.getContext());
300   Constant *ValueToRemove =
301       ConstantExpr::getPointerBitCastOrAddrSpaceCast(&Value, Int8PtrTy);
302   SmallPtrSet<Constant *, 16> InitAsSet;
303   SmallVector<Constant *, 16> Init;
304   if (GV) {
305     if (GV->hasInitializer()) {
306       auto *CA = cast<ConstantArray>(GV->getInitializer());
307       for (auto &Op : CA->operands()) {
308         Constant *C = cast_or_null<Constant>(Op);
309         if (C != ValueToRemove && InitAsSet.insert(C).second)
310           Init.push_back(C);
311       }
312     }
313     GV->eraseFromParent();
314   }
315 
316   if (Init.empty())
317     return;
318 
319   ArrayType *ATy = ArrayType::get(Int8PtrTy, Init.size());
320   GV = new llvm::GlobalVariable(M, ATy, false, GlobalValue::AppendingLinkage,
321                                 ConstantArray::get(ATy, Init),
322                                 "llvm.compiler.used");
323   GV->setSection("llvm.metadata");
324 }
325 
326 /// Attempts to extract all the embedded device images contained inside the
327 /// buffer \p Contents. The buffer is expected to contain a valid offloading
328 /// binary format.
329 Error extractOffloadFiles(StringRef Contents, StringRef Prefix,
330                           SmallVectorImpl<DeviceFile> &DeviceFiles,
331                           bool IsLibrary = false) {
332   uint64_t Offset = 0;
333   // There could be multiple offloading binaries stored at this section.
334   while (Offset < Contents.size()) {
335     std::unique_ptr<MemoryBuffer> Buffer =
336         MemoryBuffer::getMemBuffer(Contents.drop_front(Offset), "",
337                                    /*RequiresNullTerminator*/ false);
338     auto BinaryOrErr = OffloadBinary::create(*Buffer);
339     if (!BinaryOrErr)
340       return BinaryOrErr.takeError();
341     OffloadBinary &Binary = **BinaryOrErr;
342 
343     if (Binary.getVersion() != 1)
344       return createStringError(inconvertibleErrorCode(),
345                                "Incompatible device image version");
346 
347     StringRef Kind = getOffloadKindName(Binary.getOffloadKind());
348     StringRef Suffix = getImageKindName(Binary.getImageKind());
349 
350     SmallString<128> TempFile;
351     if (Error Err =
352             createOutputFile(Prefix + "-" + Kind + "-" + Binary.getTriple() +
353                                  "-" + Binary.getArch(),
354                              Suffix, TempFile))
355       return Err;
356 
357     Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
358         FileOutputBuffer::create(TempFile, Binary.getImage().size());
359     if (!OutputOrErr)
360       return OutputOrErr.takeError();
361     std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
362     std::copy(Binary.getImage().bytes_begin(), Binary.getImage().bytes_end(),
363               Output->getBufferStart());
364     if (Error E = Output->commit())
365       return E;
366 
367     DeviceFiles.emplace_back(Kind, Binary.getTriple(), Binary.getArch(),
368                              TempFile, IsLibrary);
369 
370     Offset += Binary.getSize();
371   }
372 
373   return Error::success();
374 }
375 
376 Expected<Optional<std::string>>
377 extractFromBinary(const ObjectFile &Obj,
378                   SmallVectorImpl<DeviceFile> &DeviceFiles,
379                   bool IsLibrary = false) {
380   StringRef Extension = sys::path::extension(Obj.getFileName()).drop_front();
381   StringRef Prefix = sys::path::stem(Obj.getFileName());
382   SmallVector<StringRef, 4> ToBeStripped;
383 
384   // Extract offloading binaries from sections with the name `.llvm.offloading`.
385   for (const SectionRef &Sec : Obj.sections()) {
386     Expected<StringRef> Name = Sec.getName();
387     if (!Name || !Name->equals(OFFLOAD_SECTION_MAGIC_STR))
388       continue;
389 
390     Expected<StringRef> Contents = Sec.getContents();
391     if (!Contents)
392       return Contents.takeError();
393 
394     if (Error Err =
395             extractOffloadFiles(*Contents, Prefix, DeviceFiles, IsLibrary))
396       return std::move(Err);
397 
398     ToBeStripped.push_back(*Name);
399   }
400 
401   if (ToBeStripped.empty() || !StripSections)
402     return None;
403 
404   // If the object file to strip doesn't exist we need to write it so we can
405   // pass it to llvm-strip.
406   SmallString<128> StripFile = Obj.getFileName();
407   if (!sys::fs::exists(StripFile)) {
408     SmallString<128> TempFile;
409     if (Error Err = createOutputFile(
410             sys::path::stem(StripFile),
411             sys::path::extension(StripFile).drop_front(), TempFile))
412       return std::move(Err);
413 
414     auto Contents = Obj.getMemoryBufferRef().getBuffer();
415     Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
416         FileOutputBuffer::create(TempFile, Contents.size());
417     if (!OutputOrErr)
418       return OutputOrErr.takeError();
419     std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
420     std::copy(Contents.begin(), Contents.end(), Output->getBufferStart());
421     if (Error E = Output->commit())
422       return std::move(E);
423     StripFile = TempFile;
424   }
425 
426   // We will use llvm-strip to remove the now unneeded section containing the
427   // offloading code.
428   Expected<std::string> StripPath =
429       findProgram("llvm-strip", {getMainExecutable("llvm-strip")});
430   if (!StripPath)
431     return StripPath.takeError();
432 
433   SmallString<128> TempFile;
434   if (Error Err = createOutputFile(Prefix + "-host", Extension, TempFile))
435     return std::move(Err);
436 
437   SmallVector<StringRef, 8> StripArgs;
438   StripArgs.push_back(*StripPath);
439   StripArgs.push_back("--no-strip-all");
440   StripArgs.push_back(StripFile);
441   for (auto &Section : ToBeStripped) {
442     StripArgs.push_back("--remove-section");
443     StripArgs.push_back(Section);
444   }
445   StripArgs.push_back("-o");
446   StripArgs.push_back(TempFile);
447 
448   if (Error Err = executeCommands(*StripPath, StripArgs))
449     return std::move(Err);
450 
451   return static_cast<std::string>(TempFile);
452 }
453 
454 Expected<Optional<std::string>>
455 extractFromBitcode(std::unique_ptr<MemoryBuffer> Buffer,
456                    SmallVectorImpl<DeviceFile> &DeviceFiles,
457                    bool IsLibrary = false) {
458   LLVMContext Context;
459   SMDiagnostic Err;
460   std::unique_ptr<Module> M = getLazyIRModule(std::move(Buffer), Err, Context);
461   if (!M)
462     return createStringError(inconvertibleErrorCode(),
463                              "Failed to create module");
464 
465   StringRef Extension = sys::path::extension(M->getName()).drop_front();
466   StringRef Prefix =
467       sys::path::stem(M->getName()).take_until([](char C) { return C == '-'; });
468 
469   SmallVector<GlobalVariable *, 4> ToBeDeleted;
470 
471   // Extract offloading data from globals with the `.llvm.offloading` section
472   // name.
473   for (GlobalVariable &GV : M->globals()) {
474     if (!GV.hasSection() || !GV.getSection().equals(OFFLOAD_SECTION_MAGIC_STR))
475       continue;
476 
477     auto *CDS = dyn_cast<ConstantDataSequential>(GV.getInitializer());
478     if (!CDS)
479       continue;
480 
481     StringRef Contents = CDS->getAsString();
482 
483     if (Error Err =
484             extractOffloadFiles(Contents, Prefix, DeviceFiles, IsLibrary))
485       return std::move(Err);
486 
487     ToBeDeleted.push_back(&GV);
488   }
489 
490   if (ToBeDeleted.empty() || !StripSections)
491     return None;
492 
493   // We need to materialize the lazy module before we make any changes.
494   if (Error Err = M->materializeAll())
495     return std::move(Err);
496 
497   // Remove the global from the module and write it to a new file.
498   for (GlobalVariable *GV : ToBeDeleted) {
499     removeFromCompilerUsed(*M, *GV);
500     GV->eraseFromParent();
501   }
502 
503   SmallString<128> TempFile;
504   if (Error Err = createOutputFile(Prefix + "-host", Extension, TempFile))
505     return std::move(Err);
506 
507   std::error_code EC;
508   raw_fd_ostream HostOutput(TempFile, EC, sys::fs::OF_None);
509   if (EC)
510     return createFileError(TempFile, EC);
511   WriteBitcodeToFile(*M, HostOutput);
512   return static_cast<std::string>(TempFile);
513 }
514 
515 Expected<Optional<std::string>>
516 extractFromArchive(const Archive &Library,
517                    SmallVectorImpl<DeviceFile> &DeviceFiles) {
518   bool NewMembers = false;
519   SmallVector<NewArchiveMember, 8> Members;
520 
521   // Try to extract device code from each file stored in the static archive.
522   // Save the stripped archive members to create a new host archive with the
523   // offloading code removed.
524   Error Err = Error::success();
525   for (auto Child : Library.children(Err)) {
526     auto ChildBufferRefOrErr = Child.getMemoryBufferRef();
527     if (!ChildBufferRefOrErr)
528       return ChildBufferRefOrErr.takeError();
529     std::unique_ptr<MemoryBuffer> ChildBuffer =
530         MemoryBuffer::getMemBuffer(*ChildBufferRefOrErr, false);
531 
532     auto FileOrErr = extractFromBuffer(std::move(ChildBuffer), DeviceFiles,
533                                        /*IsLibrary*/ true);
534     if (!FileOrErr)
535       return FileOrErr.takeError();
536 
537     // If we created a new stripped host file, use it to create a new archive
538     // member, otherwise use the old member.
539     if (!FileOrErr->hasValue()) {
540       Expected<NewArchiveMember> NewMember =
541           NewArchiveMember::getOldMember(Child, true);
542       if (!NewMember)
543         return NewMember.takeError();
544       Members.push_back(std::move(*NewMember));
545     } else {
546       Expected<NewArchiveMember> NewMember =
547           NewArchiveMember::getFile(**FileOrErr, true);
548       if (!NewMember)
549         return NewMember.takeError();
550       Members.push_back(std::move(*NewMember));
551       NewMembers = true;
552 
553       // We no longer need the stripped file, remove it.
554       if (std::error_code EC = sys::fs::remove(**FileOrErr))
555         return createFileError(**FileOrErr, EC);
556     }
557   }
558 
559   if (Err)
560     return std::move(Err);
561 
562   if (!NewMembers || !StripSections)
563     return None;
564 
565   // Create a new static library using the stripped host files.
566   SmallString<128> TempFile;
567   StringRef Prefix = sys::path::stem(Library.getFileName());
568   if (Error Err = createOutputFile(Prefix + "-host", "a", TempFile))
569     return std::move(Err);
570 
571   std::unique_ptr<MemoryBuffer> Buffer =
572       MemoryBuffer::getMemBuffer(Library.getMemoryBufferRef(), false);
573   if (Error Err = writeArchive(TempFile, Members, true, Library.kind(), true,
574                                Library.isThin(), std::move(Buffer)))
575     return std::move(Err);
576 
577   return static_cast<std::string>(TempFile);
578 }
579 
580 /// Extracts embedded device offloading code from a memory \p Buffer to a list
581 /// of \p DeviceFiles. If device code was extracted a new file with the embedded
582 /// device code stripped from the buffer will be returned.
583 Expected<Optional<std::string>>
584 extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
585                   SmallVectorImpl<DeviceFile> &DeviceFiles, bool IsLibrary) {
586   file_magic Type = identify_magic(Buffer->getBuffer());
587   switch (Type) {
588   case file_magic::bitcode:
589     return extractFromBitcode(std::move(Buffer), DeviceFiles, IsLibrary);
590   case file_magic::elf_relocatable:
591   case file_magic::macho_object:
592   case file_magic::coff_object: {
593     Expected<std::unique_ptr<ObjectFile>> ObjFile =
594         ObjectFile::createObjectFile(*Buffer, Type);
595     if (!ObjFile)
596       return ObjFile.takeError();
597     return extractFromBinary(*ObjFile->get(), DeviceFiles, IsLibrary);
598   }
599   case file_magic::archive: {
600     Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
601         object::Archive::create(*Buffer);
602     if (!LibFile)
603       return LibFile.takeError();
604     return extractFromArchive(*LibFile->get(), DeviceFiles);
605   }
606   default:
607     return None;
608   }
609 }
610 
611 // TODO: Move these to a separate file.
612 namespace nvptx {
613 Expected<std::string> assemble(StringRef InputFile, Triple TheTriple,
614                                StringRef Arch, bool RDC = true) {
615   // NVPTX uses the ptxas binary to create device object files.
616   Expected<std::string> PtxasPath = findProgram("ptxas", {CudaBinaryPath});
617   if (!PtxasPath)
618     return PtxasPath.takeError();
619 
620   // Create a new file to write the linked device image to.
621   SmallString<128> TempFile;
622   if (Error Err =
623           createOutputFile(sys::path::filename(ExecutableName) + "-device-" +
624                                TheTriple.getArchName() + "-" + Arch,
625                            "cubin", TempFile))
626     return std::move(Err);
627 
628   SmallVector<StringRef, 16> CmdArgs;
629   std::string Opt = "-" + OptLevel;
630   CmdArgs.push_back(*PtxasPath);
631   CmdArgs.push_back(TheTriple.isArch64Bit() ? "-m64" : "-m32");
632   if (Verbose)
633     CmdArgs.push_back("-v");
634   if (DebugInfo == DirectivesOnly && OptLevel[1] == '0')
635     CmdArgs.push_back("-lineinfo");
636   else if (DebugInfo == FullDebugInfo && OptLevel[1] == '0')
637     CmdArgs.push_back("-g");
638   for (auto &Arg : PtxasArgs)
639     CmdArgs.push_back(Arg);
640   CmdArgs.push_back("-o");
641   CmdArgs.push_back(TempFile);
642   CmdArgs.push_back(Opt);
643   CmdArgs.push_back("--gpu-name");
644   CmdArgs.push_back(Arch);
645   if (RDC)
646     CmdArgs.push_back("-c");
647 
648   CmdArgs.push_back(InputFile);
649 
650   if (Error Err = executeCommands(*PtxasPath, CmdArgs))
651     return std::move(Err);
652 
653   return static_cast<std::string>(TempFile);
654 }
655 
656 Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
657                            StringRef Arch) {
658   // NVPTX uses the nvlink binary to link device object files.
659   Expected<std::string> NvlinkPath = findProgram("nvlink", {CudaBinaryPath});
660   if (!NvlinkPath)
661     return NvlinkPath.takeError();
662 
663   // Create a new file to write the linked device image to.
664   SmallString<128> TempFile;
665   if (Error Err =
666           createOutputFile(sys::path::filename(ExecutableName) + "-device-" +
667                                TheTriple.getArchName() + "-" + Arch,
668                            "out", TempFile))
669     return std::move(Err);
670 
671   SmallVector<StringRef, 16> CmdArgs;
672   CmdArgs.push_back(*NvlinkPath);
673   CmdArgs.push_back(TheTriple.isArch64Bit() ? "-m64" : "-m32");
674   if (Verbose)
675     CmdArgs.push_back("-v");
676   if (DebugInfo != NoDebugInfo)
677     CmdArgs.push_back("-g");
678   CmdArgs.push_back("-o");
679   CmdArgs.push_back(TempFile);
680   CmdArgs.push_back("-arch");
681   CmdArgs.push_back(Arch);
682 
683   // Add extracted input files.
684   for (StringRef Input : InputFiles)
685     CmdArgs.push_back(Input);
686 
687   if (Error Err = executeCommands(*NvlinkPath, CmdArgs))
688     return std::move(Err);
689 
690   return static_cast<std::string>(TempFile);
691 }
692 } // namespace nvptx
693 namespace amdgcn {
694 Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
695                            StringRef Arch) {
696   // AMDGPU uses lld to link device object files.
697   Expected<std::string> LLDPath =
698       findProgram("lld", {getMainExecutable("lld")});
699   if (!LLDPath)
700     return LLDPath.takeError();
701 
702   // Create a new file to write the linked device image to.
703   SmallString<128> TempFile;
704   if (Error Err = createOutputFile(sys::path::filename(ExecutableName) + "-" +
705                                        TheTriple.getArchName() + "-" + Arch,
706                                    "out", TempFile))
707     return std::move(Err);
708 
709   SmallVector<StringRef, 16> CmdArgs;
710   CmdArgs.push_back(*LLDPath);
711   CmdArgs.push_back("-flavor");
712   CmdArgs.push_back("gnu");
713   CmdArgs.push_back("--no-undefined");
714   CmdArgs.push_back("-shared");
715   CmdArgs.push_back("-o");
716   CmdArgs.push_back(TempFile);
717 
718   // Add extracted input files.
719   for (StringRef Input : InputFiles)
720     CmdArgs.push_back(Input);
721 
722   if (Error Err = executeCommands(*LLDPath, CmdArgs))
723     return std::move(Err);
724 
725   return static_cast<std::string>(TempFile);
726 }
727 } // namespace amdgcn
728 
729 namespace generic {
730 
731 const char *getLDMOption(const llvm::Triple &T) {
732   switch (T.getArch()) {
733   case llvm::Triple::x86:
734     if (T.isOSIAMCU())
735       return "elf_iamcu";
736     return "elf_i386";
737   case llvm::Triple::aarch64:
738     return "aarch64linux";
739   case llvm::Triple::aarch64_be:
740     return "aarch64linuxb";
741   case llvm::Triple::ppc64:
742     return "elf64ppc";
743   case llvm::Triple::ppc64le:
744     return "elf64lppc";
745   case llvm::Triple::x86_64:
746     if (T.isX32())
747       return "elf32_x86_64";
748     return "elf_x86_64";
749   case llvm::Triple::ve:
750     return "elf64ve";
751   default:
752     return nullptr;
753   }
754 }
755 
756 Expected<std::string> link(ArrayRef<std::string> InputFiles, Triple TheTriple,
757                            StringRef Arch) {
758   // Create a new file to write the linked device image to.
759   SmallString<128> TempFile;
760   if (Error Err = createOutputFile(sys::path::filename(ExecutableName) + "-" +
761                                        TheTriple.getArchName() + "-" + Arch,
762                                    "out", TempFile))
763     return std::move(Err);
764 
765   // Use the host linker to perform generic offloading. Use the same libraries
766   // and paths as the host application does.
767   SmallVector<StringRef, 16> CmdArgs;
768   CmdArgs.push_back(LinkerUserPath);
769   CmdArgs.push_back("-m");
770   CmdArgs.push_back(getLDMOption(TheTriple));
771   CmdArgs.push_back("-shared");
772   for (auto AI = HostLinkerArgs.begin(), AE = HostLinkerArgs.end(); AI != AE;
773        ++AI) {
774     StringRef Arg = *AI;
775     if (Arg.startswith("-L"))
776       CmdArgs.push_back(Arg);
777     else if (Arg.startswith("-l"))
778       CmdArgs.push_back(Arg);
779     else if (Arg.startswith("--as-needed"))
780       CmdArgs.push_back(Arg);
781     else if (Arg.startswith("--no-as-needed"))
782       CmdArgs.push_back(Arg);
783     else if (Arg.startswith("-rpath")) {
784       CmdArgs.push_back(Arg);
785       CmdArgs.push_back(*std::next(AI));
786     } else if (Arg.startswith("-dynamic-linker")) {
787       CmdArgs.push_back(Arg);
788       CmdArgs.push_back(*std::next(AI));
789     }
790   }
791   CmdArgs.push_back("-Bsymbolic");
792   CmdArgs.push_back("-o");
793   CmdArgs.push_back(TempFile);
794 
795   // Add extracted input files.
796   for (StringRef Input : InputFiles)
797     CmdArgs.push_back(Input);
798 
799   if (Error Err = executeCommands(LinkerUserPath, CmdArgs))
800     return std::move(Err);
801 
802   return static_cast<std::string>(TempFile);
803 }
804 } // namespace generic
805 
806 Expected<std::string> linkDevice(ArrayRef<std::string> InputFiles,
807                                  Triple TheTriple, StringRef Arch) {
808   switch (TheTriple.getArch()) {
809   case Triple::nvptx:
810   case Triple::nvptx64:
811     return nvptx::link(InputFiles, TheTriple, Arch);
812   case Triple::amdgcn:
813     return amdgcn::link(InputFiles, TheTriple, Arch);
814   case Triple::x86:
815   case Triple::x86_64:
816   case Triple::aarch64:
817   case Triple::aarch64_be:
818   case Triple::ppc64:
819   case Triple::ppc64le:
820     return generic::link(InputFiles, TheTriple, Arch);
821   default:
822     return createStringError(inconvertibleErrorCode(),
823                              TheTriple.getArchName() +
824                                  " linking is not supported");
825   }
826 }
827 
828 void diagnosticHandler(const DiagnosticInfo &DI) {
829   std::string ErrStorage;
830   raw_string_ostream OS(ErrStorage);
831   DiagnosticPrinterRawOStream DP(OS);
832   DI.print(DP);
833 
834   switch (DI.getSeverity()) {
835   case DS_Error:
836     WithColor::error(errs(), LinkerExecutable) << ErrStorage << "\n";
837     break;
838   case DS_Warning:
839     WithColor::warning(errs(), LinkerExecutable) << ErrStorage << "\n";
840     break;
841   case DS_Note:
842     WithColor::note(errs(), LinkerExecutable) << ErrStorage << "\n";
843     break;
844   case DS_Remark:
845     WithColor::remark(errs()) << ErrStorage << "\n";
846     break;
847   }
848 }
849 
850 // Get the target features passed in from the driver as <triple>=<features>.
851 std::vector<std::string> getTargetFeatures(const Triple &TheTriple) {
852   std::vector<std::string> Features;
853   auto TargetAndFeatures = StringRef(TargetFeatures).split('=');
854   if (TargetAndFeatures.first != TheTriple.getTriple())
855     return Features;
856 
857   for (auto Feature : llvm::split(TargetAndFeatures.second, ','))
858     Features.push_back(Feature.str());
859   return Features;
860 }
861 
862 CodeGenOpt::Level getCGOptLevel(unsigned OptLevel) {
863   switch (OptLevel) {
864   case 0:
865     return CodeGenOpt::None;
866   case 1:
867     return CodeGenOpt::Less;
868   case 2:
869     return CodeGenOpt::Default;
870   case 3:
871     return CodeGenOpt::Aggressive;
872   }
873   llvm_unreachable("Invalid optimization level");
874 }
875 
876 template <typename ModuleHook = function_ref<bool(size_t, const Module &)>>
877 std::unique_ptr<lto::LTO> createLTO(
878     const Triple &TheTriple, StringRef Arch, bool WholeProgram,
879     ModuleHook Hook = [](size_t, const Module &) { return true; }) {
880   lto::Config Conf;
881   lto::ThinBackend Backend;
882   // TODO: Handle index-only thin-LTO
883   Backend =
884       lto::createInProcessThinBackend(llvm::heavyweight_hardware_concurrency());
885 
886   Conf.CPU = Arch.str();
887   Conf.Options = codegen::InitTargetOptionsFromCodeGenFlags(TheTriple);
888 
889   Conf.MAttrs = getTargetFeatures(TheTriple);
890   Conf.CGOptLevel = getCGOptLevel(OptLevel[1] - '0');
891   Conf.OptLevel = OptLevel[1] - '0';
892   if (Conf.OptLevel > 0)
893     Conf.UseDefaultPipeline = true;
894   Conf.DefaultTriple = TheTriple.getTriple();
895   Conf.DiagHandler = diagnosticHandler;
896 
897   Conf.PTO.LoopVectorization = Conf.OptLevel > 1;
898   Conf.PTO.SLPVectorization = Conf.OptLevel > 1;
899 
900   if (SaveTemps) {
901     auto HandleError = [&](Error Err) {
902       logAllUnhandledErrors(std::move(Err),
903                             WithColor::error(errs(), LinkerExecutable));
904       exit(1);
905     };
906     Conf.PostInternalizeModuleHook = [&](size_t, const Module &M) {
907       SmallString<128> TempFile;
908       if (Error Err = createOutputFile(sys::path::filename(ExecutableName) +
909                                            "-device-" + TheTriple.getTriple(),
910                                        "bc", TempFile))
911         HandleError(std::move(Err));
912 
913       std::error_code EC;
914       raw_fd_ostream LinkedBitcode(TempFile, EC, sys::fs::OF_None);
915       if (EC)
916         HandleError(errorCodeToError(EC));
917       WriteBitcodeToFile(M, LinkedBitcode);
918       return true;
919     };
920   }
921   Conf.PostOptModuleHook = Hook;
922   if (TheTriple.isNVPTX())
923     Conf.CGFileType = CGFT_AssemblyFile;
924   else
925     Conf.CGFileType = CGFT_ObjectFile;
926 
927   // TODO: Handle remark files
928   Conf.HasWholeProgramVisibility = WholeProgram;
929 
930   return std::make_unique<lto::LTO>(std::move(Conf), Backend);
931 }
932 
933 // Returns true if \p S is valid as a C language identifier and will be given
934 // `__start_` and `__stop_` symbols.
935 bool isValidCIdentifier(StringRef S) {
936   return !S.empty() && (isAlpha(S[0]) || S[0] == '_') &&
937          std::all_of(S.begin() + 1, S.end(),
938                      [](char C) { return C == '_' || isAlnum(C); });
939 }
940 
941 Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
942                        const Triple &TheTriple, StringRef Arch,
943                        bool &WholeProgram) {
944   SmallVector<std::unique_ptr<MemoryBuffer>, 4> SavedBuffers;
945   SmallVector<std::unique_ptr<lto::InputFile>, 4> BitcodeFiles;
946   SmallVector<std::string, 4> NewInputFiles;
947   DenseSet<StringRef> UsedInRegularObj;
948   DenseSet<StringRef> UsedInSharedLib;
949   BumpPtrAllocator Alloc;
950   StringSaver Saver(Alloc);
951 
952   // Search for bitcode files in the input and create an LTO input file. If it
953   // is not a bitcode file, scan its symbol table for symbols we need to
954   // save.
955   for (StringRef File : InputFiles) {
956     ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
957         MemoryBuffer::getFileOrSTDIN(File);
958     if (std::error_code EC = BufferOrErr.getError())
959       return createFileError(File, EC);
960     MemoryBufferRef Buffer = **BufferOrErr;
961 
962     file_magic Type = identify_magic((*BufferOrErr)->getBuffer());
963     switch (Type) {
964     case file_magic::bitcode: {
965       Expected<std::unique_ptr<lto::InputFile>> InputFileOrErr =
966           llvm::lto::InputFile::create(Buffer);
967       if (!InputFileOrErr)
968         return InputFileOrErr.takeError();
969 
970       // Save the input file and the buffer associated with its memory.
971       BitcodeFiles.push_back(std::move(*InputFileOrErr));
972       SavedBuffers.push_back(std::move(*BufferOrErr));
973       continue;
974     }
975     case file_magic::cuda_fatbinary: {
976       // Cuda fatbinaries made by Clang almost almost have an object eighty
977       // bytes from the beginning. This should be sufficient to identify the
978       // symbols.
979       Buffer = MemoryBufferRef(
980           (*BufferOrErr)->getBuffer().drop_front(FatbinaryOffset), "FatBinary");
981       LLVM_FALLTHROUGH;
982     }
983     case file_magic::elf_relocatable:
984     case file_magic::elf_shared_object:
985     case file_magic::macho_object:
986     case file_magic::coff_object: {
987       Expected<std::unique_ptr<ObjectFile>> ObjFile =
988           ObjectFile::createObjectFile(Buffer);
989       if (!ObjFile)
990         continue;
991 
992       NewInputFiles.push_back(File.str());
993       for (auto &Sym : (*ObjFile)->symbols()) {
994         Expected<StringRef> Name = Sym.getName();
995         if (!Name)
996           return Name.takeError();
997 
998         // Record if we've seen these symbols in any object or shared libraries.
999         if ((*ObjFile)->isRelocatableObject())
1000           UsedInRegularObj.insert(Saver.save(*Name));
1001         else
1002           UsedInSharedLib.insert(Saver.save(*Name));
1003       }
1004       continue;
1005     }
1006     default:
1007       continue;
1008     }
1009   }
1010 
1011   if (BitcodeFiles.empty())
1012     return Error::success();
1013 
1014   auto HandleError = [&](Error Err) {
1015     logAllUnhandledErrors(std::move(Err),
1016                           WithColor::error(errs(), LinkerExecutable));
1017     exit(1);
1018   };
1019 
1020   // LTO Module hook to output bitcode without running the backend.
1021   auto OutputBitcode = [&](size_t Task, const Module &M) {
1022     SmallString<128> TempFile;
1023     if (Error Err = createOutputFile(sys::path::filename(ExecutableName) +
1024                                          "-jit-" + TheTriple.getTriple(),
1025                                      "bc", TempFile))
1026       HandleError(std::move(Err));
1027 
1028     std::error_code EC;
1029     raw_fd_ostream LinkedBitcode(TempFile, EC, sys::fs::OF_None);
1030     if (EC)
1031       HandleError(errorCodeToError(EC));
1032     WriteBitcodeToFile(M, LinkedBitcode);
1033     NewInputFiles.push_back(static_cast<std::string>(TempFile));
1034     return false;
1035   };
1036 
1037   // We assume visibility of the whole program if every input file was bitcode.
1038   WholeProgram = BitcodeFiles.size() == InputFiles.size();
1039   auto LTOBackend =
1040       (EmbedBitcode) ? createLTO(TheTriple, Arch, WholeProgram, OutputBitcode)
1041                      : createLTO(TheTriple, Arch, WholeProgram);
1042 
1043   // We need to resolve the symbols so the LTO backend knows which symbols need
1044   // to be kept or can be internalized. This is a simplified symbol resolution
1045   // scheme to approximate the full resolution a linker would do.
1046   DenseSet<StringRef> PrevailingSymbols;
1047   for (auto &BitcodeFile : BitcodeFiles) {
1048     const auto Symbols = BitcodeFile->symbols();
1049     SmallVector<lto::SymbolResolution, 16> Resolutions(Symbols.size());
1050     size_t Idx = 0;
1051     for (auto &Sym : Symbols) {
1052       lto::SymbolResolution &Res = Resolutions[Idx++];
1053 
1054       // We will use this as the prevailing symbol definition in LTO unless
1055       // it is undefined or another definition has already been used.
1056       Res.Prevailing =
1057           !Sym.isUndefined() &&
1058           PrevailingSymbols.insert(Saver.save(Sym.getName())).second;
1059 
1060       // We need LTO to preseve the following global symbols:
1061       // 1) Symbols used in regular objects.
1062       // 2) Sections that will be given a __start/__stop symbol.
1063       // 3) Prevailing symbols that are needed visible to external libraries.
1064       Res.VisibleToRegularObj =
1065           UsedInRegularObj.contains(Sym.getName()) ||
1066           isValidCIdentifier(Sym.getSectionName()) ||
1067           (Res.Prevailing &&
1068            (Sym.getVisibility() != GlobalValue::HiddenVisibility &&
1069             !Sym.canBeOmittedFromSymbolTable()));
1070 
1071       // Identify symbols that must be exported dynamically and can be
1072       // referenced by other files.
1073       Res.ExportDynamic =
1074           Sym.getVisibility() != GlobalValue::HiddenVisibility &&
1075           (UsedInSharedLib.contains(Sym.getName()) ||
1076            !Sym.canBeOmittedFromSymbolTable());
1077 
1078       // The final definition will reside in this linkage unit if the symbol is
1079       // defined and local to the module. This only checks for bitcode files,
1080       // full assertion will require complete symbol resolution.
1081       Res.FinalDefinitionInLinkageUnit =
1082           Sym.getVisibility() != GlobalValue::DefaultVisibility &&
1083           (!Sym.isUndefined() && !Sym.isCommon());
1084 
1085       // We do not support linker redefined symbols (e.g. --wrap) for device
1086       // image linking, so the symbols will not be changed after LTO.
1087       Res.LinkerRedefined = false;
1088     }
1089 
1090     // Add the bitcode file with its resolved symbols to the LTO job.
1091     if (Error Err = LTOBackend->add(std::move(BitcodeFile), Resolutions))
1092       return Err;
1093   }
1094 
1095   // Run the LTO job to compile the bitcode.
1096   size_t MaxTasks = LTOBackend->getMaxTasks();
1097   std::vector<SmallString<128>> Files(MaxTasks);
1098   auto AddStream = [&](size_t Task) -> std::unique_ptr<CachedFileStream> {
1099     int FD = -1;
1100     auto &TempFile = Files[Task];
1101     StringRef Extension = (TheTriple.isNVPTX()) ? "s" : "o";
1102     if (Error Err = createOutputFile(sys::path::filename(ExecutableName) +
1103                                          "-device-" + TheTriple.getTriple(),
1104                                      Extension, TempFile))
1105       HandleError(std::move(Err));
1106     if (std::error_code EC = sys::fs::openFileForWrite(TempFile, FD))
1107       HandleError(errorCodeToError(EC));
1108     return std::make_unique<CachedFileStream>(
1109         std::make_unique<llvm::raw_fd_ostream>(FD, true));
1110   };
1111 
1112   if (Error Err = LTOBackend->run(AddStream))
1113     return Err;
1114 
1115   // Is we are compiling for NVPTX we need to run the assembler first.
1116   if (TheTriple.isNVPTX() && !EmbedBitcode) {
1117     for (auto &File : Files) {
1118       auto FileOrErr = nvptx::assemble(File, TheTriple, Arch, !WholeProgram);
1119       if (!FileOrErr)
1120         return FileOrErr.takeError();
1121       File = *FileOrErr;
1122     }
1123   }
1124 
1125   // Append the new inputs to the device linker input.
1126   for (auto &File : Files)
1127     NewInputFiles.push_back(static_cast<std::string>(File));
1128   InputFiles = NewInputFiles;
1129 
1130   return Error::success();
1131 }
1132 
1133 /// Runs the appropriate linking action on all the device files specified in \p
1134 /// DeviceFiles. The linked device images are returned in \p LinkedImages.
1135 Error linkDeviceFiles(ArrayRef<DeviceFile> DeviceFiles,
1136                       SmallVectorImpl<std::string> &LinkedImages) {
1137   // Get the list of inputs for a specific device.
1138   DenseMap<DeviceFile, SmallVector<std::string, 4>> LinkerInputMap;
1139   SmallVector<DeviceFile, 4> LibraryFiles;
1140   for (auto &File : DeviceFiles) {
1141     if (File.IsLibrary)
1142       LibraryFiles.push_back(File);
1143     else
1144       LinkerInputMap[File].push_back(File.Filename);
1145   }
1146 
1147   // Static libraries are loaded lazily as-needed, only add them if other files
1148   // are present.
1149   // TODO: We need to check the symbols as well, static libraries are only
1150   //       loaded if they contain symbols that are currently undefined or common
1151   //       in the symbol table.
1152   for (auto &File : LibraryFiles)
1153     if (LinkerInputMap.count(File))
1154       LinkerInputMap[File].push_back(File.Filename);
1155 
1156   // Try to link each device toolchain.
1157   for (auto &LinkerInput : LinkerInputMap) {
1158     DeviceFile &File = LinkerInput.getFirst();
1159     Triple TheTriple = Triple(File.TheTriple);
1160     bool WholeProgram = false;
1161 
1162     // Run LTO on any bitcode files and replace the input with the result.
1163     if (Error Err = linkBitcodeFiles(LinkerInput.getSecond(), TheTriple,
1164                                      File.Arch, WholeProgram))
1165       return Err;
1166 
1167     // If we are embedding bitcode for JIT, skip the final device linking.
1168     if (EmbedBitcode) {
1169       assert(!LinkerInput.getSecond().empty() && "No bitcode image to embed");
1170       LinkedImages.push_back(LinkerInput.getSecond().front());
1171       continue;
1172     }
1173 
1174     // If we performed LTO on NVPTX and had whole program visibility, we can use
1175     // CUDA in non-RDC mode.
1176     if (WholeProgram && TheTriple.isNVPTX()) {
1177       assert(!LinkerInput.getSecond().empty() && "No non-RDC image to embed");
1178       LinkedImages.push_back(LinkerInput.getSecond().front());
1179       continue;
1180     }
1181 
1182     auto ImageOrErr = linkDevice(LinkerInput.getSecond(), TheTriple, File.Arch);
1183     if (!ImageOrErr)
1184       return ImageOrErr.takeError();
1185 
1186     LinkedImages.push_back(*ImageOrErr);
1187   }
1188   return Error::success();
1189 }
1190 
1191 // Compile the module to an object file using the appropriate target machine for
1192 // the host triple.
1193 Expected<std::string> compileModule(Module &M) {
1194   std::string Msg;
1195   const Target *T = TargetRegistry::lookupTarget(M.getTargetTriple(), Msg);
1196   if (!T)
1197     return createStringError(inconvertibleErrorCode(), Msg);
1198 
1199   auto Options =
1200       codegen::InitTargetOptionsFromCodeGenFlags(Triple(M.getTargetTriple()));
1201   StringRef CPU = "";
1202   StringRef Features = "";
1203   std::unique_ptr<TargetMachine> TM(T->createTargetMachine(
1204       HostTriple, CPU, Features, Options, Reloc::PIC_, M.getCodeModel()));
1205 
1206   if (M.getDataLayout().isDefault())
1207     M.setDataLayout(TM->createDataLayout());
1208 
1209   SmallString<128> ObjectFile;
1210   int FD = -1;
1211   if (Error Err = createOutputFile(
1212           sys::path::filename(ExecutableName) + "-wrapper", "o", ObjectFile))
1213     return std::move(Err);
1214   if (std::error_code EC = sys::fs::openFileForWrite(ObjectFile, FD))
1215     return errorCodeToError(EC);
1216 
1217   auto OS = std::make_unique<llvm::raw_fd_ostream>(FD, true);
1218 
1219   legacy::PassManager CodeGenPasses;
1220   TargetLibraryInfoImpl TLII(Triple(M.getTargetTriple()));
1221   CodeGenPasses.add(new TargetLibraryInfoWrapperPass(TLII));
1222   if (TM->addPassesToEmitFile(CodeGenPasses, *OS, nullptr, CGFT_ObjectFile))
1223     return createStringError(inconvertibleErrorCode(),
1224                              "Failed to execute host backend");
1225   CodeGenPasses.run(M);
1226 
1227   return static_cast<std::string>(ObjectFile);
1228 }
1229 
1230 /// Creates the object file containing the device image and runtime registration
1231 /// code from the device images stored in \p Images.
1232 Expected<std::string> wrapDeviceImages(ArrayRef<std::string> Images) {
1233   SmallVector<std::unique_ptr<MemoryBuffer>, 4> SavedBuffers;
1234   SmallVector<ArrayRef<char>, 4> ImagesToWrap;
1235 
1236   for (StringRef ImageFilename : Images) {
1237     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ImageOrError =
1238         llvm::MemoryBuffer::getFileOrSTDIN(ImageFilename);
1239     if (std::error_code EC = ImageOrError.getError())
1240       return createFileError(ImageFilename, EC);
1241     ImagesToWrap.emplace_back((*ImageOrError)->getBufferStart(),
1242                               (*ImageOrError)->getBufferSize());
1243     SavedBuffers.emplace_back(std::move(*ImageOrError));
1244   }
1245 
1246   LLVMContext Context;
1247   Module M("offload.wrapper.module", Context);
1248   M.setTargetTriple(HostTriple);
1249   if (Error Err = wrapBinaries(M, ImagesToWrap))
1250     return std::move(Err);
1251 
1252   if (PrintWrappedModule)
1253     llvm::errs() << M;
1254 
1255   return compileModule(M);
1256 }
1257 
1258 Optional<std::string> findFile(StringRef Dir, const Twine &Name) {
1259   SmallString<128> Path;
1260   if (Dir.startswith("="))
1261     sys::path::append(Path, Sysroot, Dir.substr(1), Name);
1262   else
1263     sys::path::append(Path, Dir, Name);
1264 
1265   if (sys::fs::exists(Path))
1266     return static_cast<std::string>(Path);
1267   return None;
1268 }
1269 
1270 Optional<std::string> findFromSearchPaths(StringRef Name,
1271                                           ArrayRef<StringRef> SearchPaths) {
1272   for (StringRef Dir : SearchPaths)
1273     if (Optional<std::string> File = findFile(Dir, Name))
1274       return File;
1275   return None;
1276 }
1277 
1278 Optional<std::string> searchLibraryBaseName(StringRef Name,
1279                                             ArrayRef<StringRef> SearchPaths) {
1280   for (StringRef Dir : SearchPaths) {
1281     if (Optional<std::string> File = findFile(Dir, "lib" + Name + ".so"))
1282       return None;
1283     if (Optional<std::string> File = findFile(Dir, "lib" + Name + ".a"))
1284       return File;
1285   }
1286   return None;
1287 }
1288 
1289 /// Search for static libraries in the linker's library path given input like
1290 /// `-lfoo` or `-l:libfoo.a`.
1291 Optional<std::string> searchLibrary(StringRef Input,
1292                                     ArrayRef<StringRef> SearchPaths) {
1293   if (!Input.startswith("-l"))
1294     return None;
1295   StringRef Name = Input.drop_front(2);
1296   if (Name.startswith(":"))
1297     return findFromSearchPaths(Name.drop_front(), SearchPaths);
1298   return searchLibraryBaseName(Name, SearchPaths);
1299 }
1300 
1301 } // namespace
1302 
1303 int main(int argc, const char **argv) {
1304   InitLLVM X(argc, argv);
1305   InitializeAllTargetInfos();
1306   InitializeAllTargets();
1307   InitializeAllTargetMCs();
1308   InitializeAllAsmParsers();
1309   InitializeAllAsmPrinters();
1310 
1311   LinkerExecutable = argv[0];
1312   sys::PrintStackTraceOnErrorSignal(argv[0]);
1313   cl::SetVersionPrinter(PrintVersion);
1314   cl::HideUnrelatedOptions(ClangLinkerWrapperCategory);
1315   cl::ParseCommandLineOptions(
1316       argc, argv,
1317       "A wrapper utility over the host linker. It scans the input files for\n"
1318       "sections that require additional processing prior to linking. The tool\n"
1319       "will then transparently pass all arguments and input to the specified\n"
1320       "host linker to create the final binary.\n");
1321 
1322   if (Help) {
1323     cl::PrintHelpMessage();
1324     return EXIT_SUCCESS;
1325   }
1326 
1327   auto reportError = [argv](Error E) {
1328     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
1329     return EXIT_FAILURE;
1330   };
1331 
1332   if (!CudaPath.empty())
1333     CudaBinaryPath = CudaPath + "/bin";
1334 
1335   auto RootIt = llvm::find_if(HostLinkerArgs, [](StringRef Arg) {
1336     return Arg.startswith("--sysroot=");
1337   });
1338   if (RootIt != HostLinkerArgs.end())
1339     Sysroot = StringRef(*RootIt).split('=').second;
1340 
1341   ExecutableName = *std::next(llvm::find(HostLinkerArgs, "-o"));
1342   SmallVector<std::string, 16> LinkerArgs;
1343   for (const std::string &Arg : HostLinkerArgs)
1344     LinkerArgs.push_back(Arg);
1345 
1346   SmallVector<StringRef, 16> LibraryPaths;
1347   for (StringRef Arg : LinkerArgs) {
1348     if (Arg.startswith("-L"))
1349       LibraryPaths.push_back(Arg.drop_front(2));
1350   }
1351 
1352   // Try to extract device code from the linker input and replace the linker
1353   // input with a new file that has the device section stripped.
1354   SmallVector<DeviceFile, 4> DeviceFiles;
1355   for (std::string &Arg : LinkerArgs) {
1356     if (Arg == ExecutableName)
1357       continue;
1358 
1359     // Search for static libraries in the library link path.
1360     std::string Filename = Arg;
1361     if (Optional<std::string> Library = searchLibrary(Arg, LibraryPaths))
1362       Filename = *Library;
1363 
1364     if (sys::fs::exists(Filename) && !sys::fs::is_directory(Filename)) {
1365       ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
1366           MemoryBuffer::getFileOrSTDIN(Filename);
1367       if (std::error_code EC = BufferOrErr.getError())
1368         return reportError(createFileError(Filename, EC));
1369 
1370       auto NewFileOrErr =
1371           extractFromBuffer(std::move(*BufferOrErr), DeviceFiles);
1372 
1373       if (!NewFileOrErr)
1374         return reportError(NewFileOrErr.takeError());
1375 
1376       if (NewFileOrErr->hasValue())
1377         Arg = **NewFileOrErr;
1378     }
1379   }
1380 
1381   // Add the device bitcode libraries to the device files if any were passed in.
1382   for (StringRef LibraryStr : BitcodeLibraries)
1383     DeviceFiles.push_back(getBitcodeLibrary(LibraryStr));
1384 
1385   // Link the device images extracted from the linker input.
1386   SmallVector<std::string, 16> LinkedImages;
1387   if (Error Err = linkDeviceFiles(DeviceFiles, LinkedImages))
1388     return reportError(std::move(Err));
1389 
1390   // Wrap each linked device image into a linkable host binary and add it to the
1391   // link job's inputs.
1392   auto FileOrErr = wrapDeviceImages(LinkedImages);
1393   if (!FileOrErr)
1394     return reportError(FileOrErr.takeError());
1395   LinkerArgs.push_back(*FileOrErr);
1396 
1397   // Run the host linking job.
1398   if (Error Err = runLinker(LinkerUserPath, LinkerArgs))
1399     return reportError(std::move(Err));
1400 
1401   // Remove the temporary files created.
1402   for (const auto &TempFile : TempFiles)
1403     if (std::error_code EC = sys::fs::remove(TempFile))
1404       reportError(createFileError(TempFile, EC));
1405 
1406   return EXIT_SUCCESS;
1407 }
1408