1 //===-- clang-offload-packager/ClangOffloadPackager.cpp - file bundler ---===//
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 takes several device object files and bundles them into a single
10 // binary image using a custom binary format. This is intended to be used to
11 // embed many device files into an application to create a fat binary.
12 //
13 //===---------------------------------------------------------------------===//
14 
15 #include "clang/Basic/Version.h"
16 
17 #include "llvm/Object/Binary.h"
18 #include "llvm/Object/ObjectFile.h"
19 #include "llvm/Object/OffloadBinary.h"
20 #include "llvm/Support/CommandLine.h"
21 #include "llvm/Support/FileOutputBuffer.h"
22 #include "llvm/Support/MemoryBuffer.h"
23 #include "llvm/Support/Path.h"
24 #include "llvm/Support/Signals.h"
25 #include "llvm/Support/StringSaver.h"
26 #include "llvm/Support/WithColor.h"
27 
28 using namespace llvm;
29 using namespace llvm::object;
30 
31 static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
32 
33 static cl::OptionCategory
34     ClangOffloadPackagerCategory("clang-offload-packager options");
35 
36 static cl::opt<std::string> OutputFile("o", cl::Required,
37                                        cl::desc("Write output to <file>."),
38                                        cl::value_desc("file"),
39                                        cl::cat(ClangOffloadPackagerCategory));
40 
41 static cl::list<std::string>
42     DeviceImages("image",
43                  cl::desc("List of key and value arguments. Required keywords "
44                           "are 'file' and 'triple'."),
45                  cl::value_desc("<key>=<value>,..."),
46                  cl::cat(ClangOffloadPackagerCategory));
47 
48 static void PrintVersion(raw_ostream &OS) {
49   OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n';
50 }
51 
52 int main(int argc, const char **argv) {
53   sys::PrintStackTraceOnErrorSignal(argv[0]);
54   cl::HideUnrelatedOptions(ClangOffloadPackagerCategory);
55   cl::SetVersionPrinter(PrintVersion);
56   cl::ParseCommandLineOptions(
57       argc, argv,
58       "A utility for bundling several object files into a single binary.\n"
59       "The output binary can then be embedded into the host section table\n"
60       "to create a fatbinary containing offloading code.\n");
61 
62   if (Help) {
63     cl::PrintHelpMessage();
64     return EXIT_SUCCESS;
65   }
66 
67   auto reportError = [argv](Error E) {
68     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
69     return EXIT_FAILURE;
70   };
71 
72   SmallVector<char, 1024> BinaryData;
73   raw_svector_ostream OS(BinaryData);
74   for (StringRef Image : DeviceImages) {
75     BumpPtrAllocator Alloc;
76     StringSaver Saver(Alloc);
77 
78     StringMap<StringRef> Args;
79     for (StringRef Arg : llvm::split(Image, ",")) {
80       auto KeyAndValue = Arg.split("=");
81       if (Args.count(KeyAndValue.first))
82         Args[KeyAndValue.first] =
83             Saver.save(Args[KeyAndValue.first] + "," + KeyAndValue.second);
84       else
85         Args[KeyAndValue.first] = KeyAndValue.second;
86     }
87 
88     if (!Args.count("triple") || !Args.count("file"))
89       return reportError(createStringError(
90           inconvertibleErrorCode(),
91           "'file' and 'triple' are required image arguments"));
92 
93     OffloadBinary::OffloadingImage ImageBinary{};
94     std::unique_ptr<llvm::MemoryBuffer> DeviceImage;
95     for (const auto &KeyAndValue : Args) {
96       StringRef Key = KeyAndValue.getKey();
97       if (Key == "file") {
98         llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ObjectOrErr =
99             llvm::MemoryBuffer::getFileOrSTDIN(KeyAndValue.getValue());
100         if (std::error_code EC = ObjectOrErr.getError())
101           return reportError(errorCodeToError(EC));
102 
103         // Clang uses the '.o' suffix for LTO bitcode.
104         if (identify_magic((*ObjectOrErr)->getBuffer()) == file_magic::bitcode)
105           ImageBinary.TheImageKind = object::IMG_Bitcode;
106         else
107           ImageBinary.TheImageKind = getImageKind(
108               sys::path::extension(KeyAndValue.getValue()).drop_front());
109         ImageBinary.Image = std::move(*ObjectOrErr);
110       } else if (Key == "kind") {
111         ImageBinary.TheOffloadKind = getOffloadKind(KeyAndValue.getValue());
112       } else {
113         ImageBinary.StringData[Key] = KeyAndValue.getValue();
114       }
115     }
116     std::unique_ptr<MemoryBuffer> Buffer = OffloadBinary::write(ImageBinary);
117     if (Buffer->getBufferSize() % OffloadBinary::getAlignment() != 0)
118       return reportError(
119           createStringError(inconvertibleErrorCode(),
120                             "Offload binary has invalid size alignment"));
121     OS << Buffer->getBuffer();
122   }
123 
124   Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
125       FileOutputBuffer::create(OutputFile, BinaryData.size());
126   if (!OutputOrErr)
127     return reportError(OutputOrErr.takeError());
128   std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
129   std::copy(BinaryData.begin(), BinaryData.end(), Output->getBufferStart());
130   if (Error E = Output->commit())
131     return reportError(std::move(E));
132 }
133