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