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