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 | |
30 | using namespace llvm; |
31 | using namespace llvm::object; |
32 | |
33 | static cl::opt<bool> Help("h" , cl::desc("Alias for -help" ), cl::Hidden); |
34 | |
35 | static cl::OptionCategory |
36 | ClangOffloadPackagerCategory("clang-offload-packager options" ); |
37 | |
38 | static cl::opt<std::string> OutputFile("o" , cl::desc("Write output to <file>." ), |
39 | cl::value_desc("file" ), |
40 | cl::cat(ClangOffloadPackagerCategory)); |
41 | |
42 | static cl::opt<std::string> InputFile(cl::Positional, |
43 | cl::desc("Extract from <file>." ), |
44 | cl::value_desc("file" ), |
45 | cl::cat(ClangOffloadPackagerCategory)); |
46 | |
47 | static 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 | |
54 | static 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. |
60 | static const char *PackagerExecutable; |
61 | |
62 | static 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. |
68 | static 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 | |
82 | static 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 | |
94 | static 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 | Msg: "'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 | Msg: "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 | |
145 | static 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 *> ; |
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 | Msg: "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 | |
225 | int 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 | Msg: "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 | |