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