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