1//===-- llvm-offload-binary.cpp - offload binary management utility -------===//
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. It also
12// supports extracting these files from a known location.
13//
14//===----------------------------------------------------------------------===//
15
16#include "llvm/ADT/StringExtras.h"
17#include "llvm/BinaryFormat/Magic.h"
18#include "llvm/Object/ArchiveWriter.h"
19#include "llvm/Object/OffloadBinary.h"
20#include "llvm/Support/CommandLine.h"
21#include "llvm/Support/FileOutputBuffer.h"
22#include "llvm/Support/FileSystem.h"
23#include "llvm/Support/MemoryBuffer.h"
24#include "llvm/Support/Path.h"
25#include "llvm/Support/Signals.h"
26#include "llvm/Support/StringSaver.h"
27#include "llvm/Support/WithColor.h"
28
29using namespace llvm;
30using namespace llvm::object;
31
32static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
33
34static cl::OptionCategory OffloadBinaryCategory("llvm-offload-binary options");
35
36static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),
37 cl::value_desc("file"),
38 cl::cat(OffloadBinaryCategory));
39
40static cl::opt<std::string> InputFile(cl::Positional,
41 cl::desc("Extract from <file>."),
42 cl::value_desc("file"),
43 cl::cat(OffloadBinaryCategory));
44
45static cl::list<std::string>
46 DeviceImages("image",
47 cl::desc("List of key and value arguments. Required keywords "
48 "are 'file' and 'triple'."),
49 cl::value_desc("<key>=<value>,..."),
50 cl::cat(OffloadBinaryCategory));
51
52static cl::opt<bool>
53 CreateArchive("archive",
54 cl::desc("Write extracted files to a static archive"),
55 cl::cat(OffloadBinaryCategory));
56
57/// Path of the current binary.
58static const char *PackagerExecutable;
59
60// Get a map containing all the arguments for the image. Repeated arguments will
61// be placed in a comma separated list.
62static DenseMap<StringRef, StringRef> getImageArguments(StringRef Image,
63 StringSaver &Saver) {
64 DenseMap<StringRef, StringRef> Args;
65 for (StringRef Arg : llvm::split(Str: Image, Separator: ",")) {
66 auto [Key, Value] = Arg.split(Separator: "=");
67 auto [It, Inserted] = Args.try_emplace(Key, Args&: Value);
68 if (!Inserted)
69 It->second = Saver.save(S: It->second + "," + Value);
70 }
71
72 return Args;
73}
74
75static Error writeFile(StringRef Filename, StringRef Data) {
76 Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
77 FileOutputBuffer::create(FilePath: Filename, Size: Data.size());
78 if (!OutputOrErr)
79 return OutputOrErr.takeError();
80 std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
81 llvm::copy(Range&: Data, Out: Output->getBufferStart());
82 if (Error E = Output->commit())
83 return E;
84 return Error::success();
85}
86
87static Error bundleImages() {
88 SmallVector<char, 1024> BinaryData;
89 raw_svector_ostream OS(BinaryData);
90 for (StringRef Image : DeviceImages) {
91 BumpPtrAllocator Alloc;
92 StringSaver Saver(Alloc);
93 DenseMap<StringRef, StringRef> Args = getImageArguments(Image, Saver);
94
95 if (!Args.count(Val: "file"))
96 return createStringError(EC: inconvertibleErrorCode(),
97 S: "'file' is a required image arguments");
98
99 // Permit using multiple instances of `file` in a single string.
100 for (auto &File : llvm::split(Str: Args["file"], Separator: ",")) {
101 OffloadBinary::OffloadingImage ImageBinary{};
102
103 llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ObjectOrErr =
104 llvm::MemoryBuffer::getFileOrSTDIN(Filename: File);
105 if (std::error_code EC = ObjectOrErr.getError())
106 return errorCodeToError(EC);
107
108 // Clang uses the '.o' suffix for LTO bitcode.
109 if (identify_magic(magic: (*ObjectOrErr)->getBuffer()) == file_magic::bitcode)
110 ImageBinary.TheImageKind = object::IMG_Bitcode;
111 else if (sys::path::has_extension(path: File))
112 ImageBinary.TheImageKind =
113 getImageKind(Name: sys::path::extension(path: File).drop_front());
114 else
115 ImageBinary.TheImageKind = IMG_None;
116 ImageBinary.Image = std::move(*ObjectOrErr);
117 for (const auto &[Key, Value] : Args) {
118 if (Key == "kind") {
119 ImageBinary.TheOffloadKind = getOffloadKind(Name: Value);
120 } else if (Key != "file") {
121 ImageBinary.StringData[Key] = Value;
122 }
123 }
124 llvm::SmallString<0> Buffer = OffloadBinary::write(ImageBinary);
125 if (Buffer.size() % OffloadBinary::getAlignment() != 0)
126 return createStringError(EC: inconvertibleErrorCode(),
127 S: "Offload binary has invalid size alignment");
128 OS << Buffer;
129 }
130 }
131
132 if (Error E = writeFile(Filename: OutputFile,
133 Data: StringRef(BinaryData.begin(), BinaryData.size())))
134 return E;
135 return Error::success();
136}
137
138static Error unbundleImages() {
139 ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
140 MemoryBuffer::getFileOrSTDIN(Filename: InputFile);
141 if (std::error_code EC = BufferOrErr.getError())
142 return createFileError(F: InputFile, EC);
143 std::unique_ptr<MemoryBuffer> Buffer = std::move(*BufferOrErr);
144
145 // This data can be misaligned if extracted from an archive.
146 if (!isAddrAligned(Lhs: Align(OffloadBinary::getAlignment()),
147 Addr: Buffer->getBufferStart()))
148 Buffer = MemoryBuffer::getMemBufferCopy(InputData: Buffer->getBuffer(),
149 BufferName: Buffer->getBufferIdentifier());
150
151 SmallVector<OffloadFile> Binaries;
152 if (Error Err = extractOffloadBinaries(Buffer: *Buffer, Binaries))
153 return Err;
154
155 // Try to extract each device image specified by the user from the input file.
156 for (StringRef Image : DeviceImages) {
157 BumpPtrAllocator Alloc;
158 StringSaver Saver(Alloc);
159 auto Args = getImageArguments(Image, Saver);
160
161 SmallVector<const OffloadBinary *> Extracted;
162 for (const OffloadFile &File : Binaries) {
163 const auto *Binary = File.getBinary();
164 // We handle the 'file' and 'kind' identifiers differently.
165 bool Match = llvm::all_of(Range&: Args, P: [&](auto &Arg) {
166 const auto [Key, Value] = Arg;
167 if (Key == "file")
168 return true;
169 if (Key == "kind")
170 return Binary->getOffloadKind() == getOffloadKind(Value);
171 return Binary->getString(Key) == Value;
172 });
173 if (Match)
174 Extracted.push_back(Elt: Binary);
175 }
176
177 if (Extracted.empty())
178 continue;
179
180 if (CreateArchive) {
181 if (!Args.count(Val: "file"))
182 return createStringError(EC: inconvertibleErrorCode(),
183 S: "Image must have a 'file' argument.");
184
185 SmallVector<NewArchiveMember> Members;
186 for (const OffloadBinary *Binary : Extracted)
187 Members.emplace_back(Args: MemoryBufferRef(
188 Binary->getImage(),
189 Binary->getMemoryBufferRef().getBufferIdentifier()));
190
191 if (Error E = writeArchive(
192 ArcName: Args["file"], NewMembers: Members, WriteSymtab: SymtabWritingMode::NormalSymtab,
193 Kind: Archive::getDefaultKind(), Deterministic: true, Thin: false, OldArchiveBuf: nullptr))
194 return E;
195 } else if (auto It = Args.find(Val: "file"); It != Args.end()) {
196 if (Extracted.size() > 1)
197 WithColor::warning(OS&: errs(), Prefix: PackagerExecutable)
198 << "Multiple inputs match to a single file, '" << It->second
199 << "'\n";
200 if (Error E = writeFile(Filename: It->second, Data: Extracted.back()->getImage()))
201 return E;
202 } else {
203 uint64_t Idx = 0;
204 for (const OffloadBinary *Binary : Extracted) {
205 StringRef Filename =
206 Saver.save(S: sys::path::stem(path: InputFile) + "-" + Binary->getTriple() +
207 "-" + Binary->getArch() + "." + std::to_string(val: Idx++) +
208 "." + getImageKindName(Name: Binary->getImageKind()));
209 if (Error E = writeFile(Filename, Data: Binary->getImage()))
210 return E;
211 }
212 }
213 }
214
215 return Error::success();
216}
217
218int main(int argc, const char **argv) {
219 sys::PrintStackTraceOnErrorSignal(Argv0: argv[0]);
220 cl::HideUnrelatedOptions(Category&: OffloadBinaryCategory);
221 cl::ParseCommandLineOptions(
222 argc, argv,
223 Overview: "A utility for bundling several object files into a single binary.\n"
224 "The output binary can then be embedded into the host section table\n"
225 "to create a fatbinary containing offloading code.\n");
226
227 if (sys::path::stem(path: argv[0]).ends_with(Suffix: "clang-offload-packager"))
228 WithColor::warning(OS&: errs(), Prefix: PackagerExecutable)
229 << "'clang-offload-packager' is deprecated. Use 'llvm-offload-binary' "
230 "instead.\n";
231
232 if (Help || (OutputFile.empty() && InputFile.empty())) {
233 cl::PrintHelpMessage();
234 return EXIT_SUCCESS;
235 }
236
237 PackagerExecutable = argv[0];
238 auto reportError = [argv](Error E) {
239 logAllUnhandledErrors(E: std::move(E), OS&: WithColor::error(OS&: errs(), Prefix: argv[0]));
240 return EXIT_FAILURE;
241 };
242
243 if (!InputFile.empty() && !OutputFile.empty())
244 return reportError(
245 createStringError(EC: inconvertibleErrorCode(),
246 S: "Packaging to an output file and extracting from an "
247 "input file are mutually exclusive."));
248
249 if (!OutputFile.empty()) {
250 if (Error Err = bundleImages())
251 return reportError(std::move(Err));
252 } else if (!InputFile.empty()) {
253 if (Error Err = unbundleImages())
254 return reportError(std::move(Err));
255 }
256
257 return EXIT_SUCCESS;
258}
259