1	//===- OffloadBundler.cpp - File Bundling and Unbundling ------------------===//
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	/// \file
10	/// This file implements an offload bundling API that bundles different files
11	/// that relate with the same source code but different targets into a single
12	/// one. Also the implements the opposite functionality, i.e. unbundle files
13	/// previous created by this API.
14	///
15	//===----------------------------------------------------------------------===//
16
17	#include "clang/Driver/OffloadBundler.h"
18	#include "clang/Basic/Cuda.h"
19	#include "clang/Basic/TargetID.h"
20	#include "llvm/ADT/ArrayRef.h"
21	#include "llvm/ADT/SmallString.h"
22	#include "llvm/ADT/SmallVector.h"
23	#include "llvm/ADT/StringExtras.h"
24	#include "llvm/ADT/StringMap.h"
25	#include "llvm/ADT/StringRef.h"
26	#include "llvm/BinaryFormat/Magic.h"
27	#include "llvm/Object/Archive.h"
28	#include "llvm/Object/ArchiveWriter.h"
29	#include "llvm/Object/Binary.h"
30	#include "llvm/Object/ObjectFile.h"
31	#include "llvm/Support/Casting.h"
32	#include "llvm/Support/Compiler.h"
33	#include "llvm/Support/Compression.h"
34	#include "llvm/Support/Debug.h"
35	#include "llvm/Support/EndianStream.h"
36	#include "llvm/Support/Errc.h"
37	#include "llvm/Support/Error.h"
38	#include "llvm/Support/ErrorOr.h"
39	#include "llvm/Support/FileSystem.h"
40	#include "llvm/Support/MD5.h"
41	#include "llvm/Support/ManagedStatic.h"
42	#include "llvm/Support/MemoryBuffer.h"
43	#include "llvm/Support/Path.h"
44	#include "llvm/Support/Program.h"
45	#include "llvm/Support/Signals.h"
46	#include "llvm/Support/StringSaver.h"
47	#include "llvm/Support/Timer.h"
48	#include "llvm/Support/WithColor.h"
49	#include "llvm/Support/raw_ostream.h"
50	#include "llvm/TargetParser/Host.h"
51	#include "llvm/TargetParser/Triple.h"
52	#include <algorithm>
53	#include <cassert>
54	#include <cstddef>
55	#include <cstdint>
56	#include <forward_list>
57	#include <llvm/Support/Process.h>
58	#include <memory>
59	#include <set>
60	#include <string>
61	#include <system_error>
62	#include <utility>
63
64	using namespace llvm;
65	using namespace llvm::object;
66	using namespace clang;
67
68	namespace {
69	struct CreateClangOffloadBundlerTimerGroup {
70	static void *call() {
71	return new TimerGroup("Clang Offload Bundler Timer Group",
72	"Timer group for clang offload bundler");
73	}
74	};
75	} // namespace
76	static llvm::ManagedStatic<llvm::TimerGroup,
77	CreateClangOffloadBundlerTimerGroup>
78	ClangOffloadBundlerTimerGroup;
79
80	/// Magic string that marks the existence of offloading data.
81	#define OFFLOAD_BUNDLER_MAGIC_STR "__CLANG_OFFLOAD_BUNDLE__"
82
83	OffloadTargetInfo::OffloadTargetInfo(const StringRef Target,
84	const OffloadBundlerConfig &BC)
85	: BundlerConfig(BC) {
86
87	// <kind>-<triple>[-<target id>[:target features]]
88	// <triple> := <arch>-<vendor>-<os>-<env>
89	SmallVector<StringRef, 6> Components;
90	Target.split(A&: Components, Separator: '-', /*MaxSplit=*/5);
91	assert((Components.size() == 5 || Components.size() == 6) &&
92	"malformed target string");
93
94	StringRef TargetIdWithFeature =
95	Components.size() == 6 ? Components.back() : "";
96	StringRef TargetId = TargetIdWithFeature.split(Separator: ':').first;
97	if (!TargetId.empty() &&
98	clang::StringToOffloadArch(S: TargetId) != clang::OffloadArch::UNKNOWN)
99	this->TargetID = TargetIdWithFeature;
100	else
101	this->TargetID = "";
102
103	this->OffloadKind = Components.front();
104	ArrayRef<StringRef> TripleSlice{&Components[1], /*length=*/4};
105	llvm::Triple T = llvm::Triple(llvm::join(R&: TripleSlice, Separator: "-"));
106	this->Triple = llvm::Triple(T.getArchName(), T.getVendorName(), T.getOSName(),
107	T.getEnvironmentName());
108	}
109
110	bool OffloadTargetInfo::hasHostKind() const {
111	return this->OffloadKind == "host";
112	}
113
114	bool OffloadTargetInfo::isOffloadKindValid() const {
115	return OffloadKind == "host" || OffloadKind == "openmp" ||
116	OffloadKind == "hip" || OffloadKind == "hipv4";
117	}
118
119	bool OffloadTargetInfo::isOffloadKindCompatible(
120	const StringRef TargetOffloadKind) const {
121	if ((OffloadKind == TargetOffloadKind) ||
122	(OffloadKind == "hip" && TargetOffloadKind == "hipv4") ||
123	(OffloadKind == "hipv4" && TargetOffloadKind == "hip"))
124	return true;
125
126	if (BundlerConfig.HipOpenmpCompatible) {
127	bool HIPCompatibleWithOpenMP = OffloadKind.starts_with_insensitive(Prefix: "hip") &&
128	TargetOffloadKind == "openmp";
129	bool OpenMPCompatibleWithHIP =
130	OffloadKind == "openmp" &&
131	TargetOffloadKind.starts_with_insensitive(Prefix: "hip");
132	return HIPCompatibleWithOpenMP || OpenMPCompatibleWithHIP;
133	}
134	return false;
135	}
136
137	bool OffloadTargetInfo::isTripleValid() const {
138	return !Triple.str().empty() && Triple.getArch() != Triple::UnknownArch;
139	}
140
141	bool OffloadTargetInfo::operator==(const OffloadTargetInfo &Target) const {
142	return OffloadKind == Target.OffloadKind &&
143	Triple.isCompatibleWith(Other: Target.Triple) && TargetID == Target.TargetID;
144	}
145
146	std::string OffloadTargetInfo::str() const {
147	std::string NormalizedTriple;
148	// Unfortunately we need some special sauce for AMDHSA because all the runtime
149	// assumes the triple to be "amdgcn/spirv64-amd-amdhsa-" (empty environment)
150	// instead of "amdgcn/spirv64-amd-amdhsa-unknown". It's gonna be very tricky
151	// to patch different layers of runtime.
152	if (Triple.getOS() == Triple::OSType::AMDHSA) {
153	NormalizedTriple = Triple.normalize(Form: Triple::CanonicalForm::THREE_IDENT);
154	NormalizedTriple.push_back(c: '-');
155	} else {
156	NormalizedTriple = Triple.normalize(Form: Triple::CanonicalForm::FOUR_IDENT);
157	}
158	return Twine(OffloadKind + "-" + NormalizedTriple + "-" + TargetID).str();
159	}
160
161	static StringRef getDeviceFileExtension(StringRef Device,
162	StringRef BundleFileName) {
163	if (Device.contains(Other: "gfx"))
164	return ".bc";
165	if (Device.contains(Other: "sm_"))
166	return ".cubin";
167	return sys::path::extension(path: BundleFileName);
168	}
169
170	static std::string getDeviceLibraryFileName(StringRef BundleFileName,
171	StringRef Device) {
172	StringRef LibName = sys::path::stem(path: BundleFileName);
173	StringRef Extension = getDeviceFileExtension(Device, BundleFileName);
174
175	std::string Result;
176	Result += LibName;
177	Result += Extension;
178	return Result;
179	}
180
181	namespace {
182	/// Generic file handler interface.
183	class FileHandler {
184	public:
185	struct BundleInfo {
186	StringRef BundleID;
187	};
188
189	FileHandler() {}
190
191	virtual ~FileHandler() {}
192
193	/// Update the file handler with information from the header of the bundled
194	/// file.
195	virtual Error ReadHeader(MemoryBuffer &Input) = 0;
196
197	/// Read the marker of the next bundled to be read in the file. The bundle
198	/// name is returned if there is one in the file, or `std::nullopt` if there
199	/// are no more bundles to be read.
200	virtual Expected<std::optional<StringRef>>
201	ReadBundleStart(MemoryBuffer &Input) = 0;
202
203	/// Read the marker that closes the current bundle.
204	virtual Error ReadBundleEnd(MemoryBuffer &Input) = 0;
205
206	/// Read the current bundle and write the result into the stream \a OS.
207	virtual Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) = 0;
208
209	/// Write the header of the bundled file to \a OS based on the information
210	/// gathered from \a Inputs.
211	virtual Error WriteHeader(raw_ostream &OS,
212	ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) = 0;
213
214	/// Write the marker that initiates a bundle for the triple \a TargetTriple to
215	/// \a OS.
216	virtual Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) = 0;
217
218	/// Write the marker that closes a bundle for the triple \a TargetTriple to \a
219	/// OS.
220	virtual Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) = 0;
221
222	/// Write the bundle from \a Input into \a OS.
223	virtual Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) = 0;
224
225	/// Finalize output file.
226	virtual Error finalizeOutputFile() { return Error::success(); }
227
228	/// List bundle IDs in \a Input.
229	virtual Error listBundleIDs(MemoryBuffer &Input) {
230	if (Error Err = ReadHeader(Input))
231	return Err;
232	return forEachBundle(Input, Func: [&](const BundleInfo &Info) -> Error {
233	llvm::outs() << Info.BundleID << '\n';
234	Error Err = listBundleIDsCallback(Input, Info);
235	if (Err)
236	return Err;
237	return Error::success();
238	});
239	}
240
241	/// Get bundle IDs in \a Input in \a BundleIds.
242	virtual Error getBundleIDs(MemoryBuffer &Input,
243	std::set<StringRef> &BundleIds) {
244	if (Error Err = ReadHeader(Input))
245	return Err;
246	return forEachBundle(Input, Func: [&](const BundleInfo &Info) -> Error {
247	BundleIds.insert(x: Info.BundleID);
248	Error Err = listBundleIDsCallback(Input, Info);
249	if (Err)
250	return Err;
251	return Error::success();
252	});
253	}
254
255	/// For each bundle in \a Input, do \a Func.
256	Error forEachBundle(MemoryBuffer &Input,
257	std::function<Error(const BundleInfo &)> Func) {
258	while (true) {
259	Expected<std::optional<StringRef>> CurTripleOrErr =
260	ReadBundleStart(Input);
261	if (!CurTripleOrErr)
262	return CurTripleOrErr.takeError();
263
264	// No more bundles.
265	if (!*CurTripleOrErr)
266	break;
267
268	StringRef CurTriple = **CurTripleOrErr;
269	assert(!CurTriple.empty());
270
271	BundleInfo Info{.BundleID: CurTriple};
272	if (Error Err = Func(Info))
273	return Err;
274	}
275	return Error::success();
276	}
277
278	protected:
279	virtual Error listBundleIDsCallback(MemoryBuffer &Input,
280	const BundleInfo &Info) {
281	return Error::success();
282	}
283	};
284
285	/// Handler for binary files. The bundled file will have the following format
286	/// (all integers are stored in little-endian format):
287	///
288	/// "OFFLOAD_BUNDLER_MAGIC_STR" (ASCII encoding of the string)
289	///
290	/// NumberOfOffloadBundles (8-byte integer)
291	///
292	/// OffsetOfBundle1 (8-byte integer)
293	/// SizeOfBundle1 (8-byte integer)
294	/// NumberOfBytesInTripleOfBundle1 (8-byte integer)
295	/// TripleOfBundle1 (byte length defined before)
296	///
297	/// ...
298	///
299	/// OffsetOfBundleN (8-byte integer)
300	/// SizeOfBundleN (8-byte integer)
301	/// NumberOfBytesInTripleOfBundleN (8-byte integer)
302	/// TripleOfBundleN (byte length defined before)
303	///
304	/// Bundle1
305	/// ...
306	/// BundleN
307
308	/// Read 8-byte integers from a buffer in little-endian format.
309	static uint64_t Read8byteIntegerFromBuffer(StringRef Buffer, size_t pos) {
310	return llvm::support::endian::read64le(P: Buffer.data() + pos);
311	}
312
313	/// Write 8-byte integers to a buffer in little-endian format.
314	static void Write8byteIntegerToBuffer(raw_ostream &OS, uint64_t Val) {
315	llvm::support::endian::write(os&: OS, value: Val, endian: llvm::endianness::little);
316	}
317
318	class BinaryFileHandler final : public FileHandler {
319	/// Information about the bundles extracted from the header.
320	struct BinaryBundleInfo final : public BundleInfo {
321	/// Size of the bundle.
322	uint64_t Size = 0u;
323	/// Offset at which the bundle starts in the bundled file.
324	uint64_t Offset = 0u;
325
326	BinaryBundleInfo() {}
327	BinaryBundleInfo(uint64_t Size, uint64_t Offset)
328	: Size(Size), Offset(Offset) {}
329	};
330
331	/// Map between a triple and the corresponding bundle information.
332	StringMap<BinaryBundleInfo> BundlesInfo;
333
334	/// Iterator for the bundle information that is being read.
335	StringMap<BinaryBundleInfo>::iterator CurBundleInfo;
336	StringMap<BinaryBundleInfo>::iterator NextBundleInfo;
337
338	/// Current bundle target to be written.
339	std::string CurWriteBundleTarget;
340
341	/// Configuration options and arrays for this bundler job
342	const OffloadBundlerConfig &BundlerConfig;
343
344	public:
345	// TODO: Add error checking from ClangOffloadBundler.cpp
346	BinaryFileHandler(const OffloadBundlerConfig &BC) : BundlerConfig(BC) {}
347
348	~BinaryFileHandler() final {}
349
350	Error ReadHeader(MemoryBuffer &Input) final {
351	StringRef FC = Input.getBuffer();
352
353	// Initialize the current bundle with the end of the container.
354	CurBundleInfo = BundlesInfo.end();
355
356	// Check if buffer is smaller than magic string.
357	size_t ReadChars = sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1;
358	if (ReadChars > FC.size())
359	return Error::success();
360
361	// Check if no magic was found.
362	if (llvm::identify_magic(magic: FC) != llvm::file_magic::offload_bundle)
363	return Error::success();
364
365	// Read number of bundles.
366	if (ReadChars + 8 > FC.size())
367	return Error::success();
368
369	uint64_t NumberOfBundles = Read8byteIntegerFromBuffer(Buffer: FC, pos: ReadChars);
370	ReadChars += 8;
371
372	// Read bundle offsets, sizes and triples.
373	for (uint64_t i = 0; i < NumberOfBundles; ++i) {
374
375	// Read offset.
376	if (ReadChars + 8 > FC.size())
377	return Error::success();
378
379	uint64_t Offset = Read8byteIntegerFromBuffer(Buffer: FC, pos: ReadChars);
380	ReadChars += 8;
381
382	// Read size.
383	if (ReadChars + 8 > FC.size())
384	return Error::success();
385
386	uint64_t Size = Read8byteIntegerFromBuffer(Buffer: FC, pos: ReadChars);
387	ReadChars += 8;
388
389	// Read triple size.
390	if (ReadChars + 8 > FC.size())
391	return Error::success();
392
393	uint64_t TripleSize = Read8byteIntegerFromBuffer(Buffer: FC, pos: ReadChars);
394	ReadChars += 8;
395
396	// Read triple.
397	if (ReadChars + TripleSize > FC.size())
398	return Error::success();
399
400	StringRef Triple(&FC.data()[ReadChars], TripleSize);
401	ReadChars += TripleSize;
402
403	// Check if the offset and size make sense.
404	if (!Offset || Offset + Size > FC.size())
405	return Error::success();
406
407	assert(!BundlesInfo.contains(Triple) && "Triple is duplicated??");
408	BundlesInfo[Triple] = BinaryBundleInfo(Size, Offset);
409	}
410	// Set the iterator to where we will start to read.
411	CurBundleInfo = BundlesInfo.end();
412	NextBundleInfo = BundlesInfo.begin();
413	return Error::success();
414	}
415
416	Expected<std::optional<StringRef>>
417	ReadBundleStart(MemoryBuffer &Input) final {
418	if (NextBundleInfo == BundlesInfo.end())
419	return std::nullopt;
420	CurBundleInfo = NextBundleInfo++;
421	return CurBundleInfo->first();
422	}
423
424	Error ReadBundleEnd(MemoryBuffer &Input) final {
425	assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
426	return Error::success();
427	}
428
429	Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
430	assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
431	StringRef FC = Input.getBuffer();
432	OS.write(Ptr: FC.data() + CurBundleInfo->second.Offset,
433	Size: CurBundleInfo->second.Size);
434	return Error::success();
435	}
436
437	Error WriteHeader(raw_ostream &OS,
438	ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
439
440	// Compute size of the header.
441	uint64_t HeaderSize = 0;
442
443	HeaderSize += sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1;
444	HeaderSize += 8; // Number of Bundles
445
446	for (auto &T : BundlerConfig.TargetNames) {
447	HeaderSize += 3 * 8; // Bundle offset, Size of bundle and size of triple.
448	HeaderSize += T.size(); // The triple.
449	}
450
451	// Write to the buffer the header.
452	OS << OFFLOAD_BUNDLER_MAGIC_STR;
453
454	Write8byteIntegerToBuffer(OS, Val: BundlerConfig.TargetNames.size());
455
456	unsigned Idx = 0;
457	for (auto &T : BundlerConfig.TargetNames) {
458	MemoryBuffer &MB = *Inputs[Idx++];
459	HeaderSize = alignTo(Value: HeaderSize, Align: BundlerConfig.BundleAlignment);
460	// Bundle offset.
461	Write8byteIntegerToBuffer(OS, Val: HeaderSize);
462	// Size of the bundle (adds to the next bundle's offset)
463	Write8byteIntegerToBuffer(OS, Val: MB.getBufferSize());
464	BundlesInfo[T] = BinaryBundleInfo(MB.getBufferSize(), HeaderSize);
465	HeaderSize += MB.getBufferSize();
466	// Size of the triple
467	Write8byteIntegerToBuffer(OS, Val: T.size());
468	// Triple
469	OS << T;
470	}
471	return Error::success();
472	}
473
474	Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) final {
475	CurWriteBundleTarget = TargetTriple.str();
476	return Error::success();
477	}
478
479	Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) final {
480	return Error::success();
481	}
482
483	Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) final {
484	auto BI = BundlesInfo[CurWriteBundleTarget];
485
486	// Pad with 0 to reach specified offset.
487	size_t CurrentPos = OS.tell();
488	size_t PaddingSize = BI.Offset > CurrentPos ? BI.Offset - CurrentPos : 0;
489	for (size_t I = 0; I < PaddingSize; ++I)
490	OS.write(C: '\0');
491	assert(OS.tell() == BI.Offset);
492
493	OS.write(Ptr: Input.getBufferStart(), Size: Input.getBufferSize());
494
495	return Error::success();
496	}
497	};
498
499	// This class implements a list of temporary files that are removed upon
500	// object destruction.
501	class TempFileHandlerRAII {
502	public:
503	~TempFileHandlerRAII() {
504	for (const auto &File : Files)
505	sys::fs::remove(path: File);
506	}
507
508	// Creates temporary file with given contents.
509	Expected<StringRef> Create(std::optional<ArrayRef<char>> Contents) {
510	SmallString<128u> File;
511	if (std::error_code EC =
512	sys::fs::createTemporaryFile(Prefix: "clang-offload-bundler", Suffix: "tmp", ResultPath&: File))
513	return createFileError(F: File, EC);
514	Files.push_front(val: File);
515
516	if (Contents) {
517	std::error_code EC;
518	raw_fd_ostream OS(File, EC);
519	if (EC)
520	return createFileError(F: File, EC);
521	OS.write(Ptr: Contents->data(), Size: Contents->size());
522	}
523	return Files.front().str();
524	}
525
526	private:
527	std::forward_list<SmallString<128u>> Files;
528	};
529
530	/// Handler for object files. The bundles are organized by sections with a
531	/// designated name.
532	///
533	/// To unbundle, we just copy the contents of the designated section.
534	class ObjectFileHandler final : public FileHandler {
535
536	/// The object file we are currently dealing with.
537	std::unique_ptr<ObjectFile> Obj;
538
539	/// Return the input file contents.
540	StringRef getInputFileContents() const { return Obj->getData(); }
541
542	/// Return bundle name (<kind>-<triple>) if the provided section is an offload
543	/// section.
544	static Expected<std::optional<StringRef>>
545	IsOffloadSection(SectionRef CurSection) {
546	Expected<StringRef> NameOrErr = CurSection.getName();
547	if (!NameOrErr)
548	return NameOrErr.takeError();
549
550	// If it does not start with the reserved suffix, just skip this section.
551	if (llvm::identify_magic(magic: *NameOrErr) != llvm::file_magic::offload_bundle)
552	return std::nullopt;
553
554	// Return the triple that is right after the reserved prefix.
555	return NameOrErr->substr(Start: sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1);
556	}
557
558	/// Total number of inputs.
559	unsigned NumberOfInputs = 0;
560
561	/// Total number of processed inputs, i.e, inputs that were already
562	/// read from the buffers.
563	unsigned NumberOfProcessedInputs = 0;
564
565	/// Iterator of the current and next section.
566	section_iterator CurrentSection;
567	section_iterator NextSection;
568
569	/// Configuration options and arrays for this bundler job
570	const OffloadBundlerConfig &BundlerConfig;
571
572	public:
573	// TODO: Add error checking from ClangOffloadBundler.cpp
574	ObjectFileHandler(std::unique_ptr<ObjectFile> ObjIn,
575	const OffloadBundlerConfig &BC)
576	: Obj(std::move(ObjIn)), CurrentSection(Obj->section_begin()),
577	NextSection(Obj->section_begin()), BundlerConfig(BC) {}
578
579	~ObjectFileHandler() final {}
580
581	Error ReadHeader(MemoryBuffer &Input) final { return Error::success(); }
582
583	Expected<std::optional<StringRef>>
584	ReadBundleStart(MemoryBuffer &Input) final {
585	while (NextSection != Obj->section_end()) {
586	CurrentSection = NextSection;
587	++NextSection;
588
589	// Check if the current section name starts with the reserved prefix. If
590	// so, return the triple.
591	Expected<std::optional<StringRef>> TripleOrErr =
592	IsOffloadSection(CurSection: *CurrentSection);
593	if (!TripleOrErr)
594	return TripleOrErr.takeError();
595	if (*TripleOrErr)
596	return **TripleOrErr;
597	}
598	return std::nullopt;
599	}
600
601	Error ReadBundleEnd(MemoryBuffer &Input) final { return Error::success(); }
602
603	Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
604	Expected<StringRef> ContentOrErr = CurrentSection->getContents();
605	if (!ContentOrErr)
606	return ContentOrErr.takeError();
607	StringRef Content = *ContentOrErr;
608
609	// Copy fat object contents to the output when extracting host bundle.
610	std::string ModifiedContent;
611	if (Content.size() == 1u && Content.front() == 0) {
612	auto HostBundleOrErr = getHostBundle(
613	Input: StringRef(Input.getBufferStart(), Input.getBufferSize()));
614	if (!HostBundleOrErr)
615	return HostBundleOrErr.takeError();
616
617	ModifiedContent = std::move(*HostBundleOrErr);
618	Content = ModifiedContent;
619	}
620
621	OS.write(Ptr: Content.data(), Size: Content.size());
622	return Error::success();
623	}
624
625	Error WriteHeader(raw_ostream &OS,
626	ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
627	assert(BundlerConfig.HostInputIndex != ~0u &&
628	"Host input index not defined.");
629
630	// Record number of inputs.
631	NumberOfInputs = Inputs.size();
632	return Error::success();
633	}
634
635	Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) final {
636	++NumberOfProcessedInputs;
637	return Error::success();
638	}
639
640	Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) final {
641	return Error::success();
642	}
643
644	Error finalizeOutputFile() final {
645	assert(NumberOfProcessedInputs <= NumberOfInputs &&
646	"Processing more inputs that actually exist!");
647	assert(BundlerConfig.HostInputIndex != ~0u &&
648	"Host input index not defined.");
649
650	// If this is not the last output, we don't have to do anything.
651	if (NumberOfProcessedInputs != NumberOfInputs)
652	return Error::success();
653
654	// We will use llvm-objcopy to add target objects sections to the output
655	// fat object. These sections should have 'exclude' flag set which tells
656	// link editor to remove them from linker inputs when linking executable or
657	// shared library.
658
659	assert(BundlerConfig.ObjcopyPath != "" &&
660	"llvm-objcopy path not specified");
661
662	// Temporary files that need to be removed.
663	TempFileHandlerRAII TempFiles;
664
665	// Compose llvm-objcopy command line for add target objects' sections with
666	// appropriate flags.
667	BumpPtrAllocator Alloc;
668	StringSaver SS{Alloc};
669	SmallVector<StringRef, 8u> ObjcopyArgs{"llvm-objcopy"};
670
671	for (unsigned I = 0; I < NumberOfInputs; ++I) {
672	StringRef InputFile = BundlerConfig.InputFileNames[I];
673	if (I == BundlerConfig.HostInputIndex) {
674	// Special handling for the host bundle. We do not need to add a
675	// standard bundle for the host object since we are going to use fat
676	// object as a host object. Therefore use dummy contents (one zero byte)
677	// when creating section for the host bundle.
678	Expected<StringRef> TempFileOrErr = TempFiles.Create(Contents: ArrayRef<char>(0));
679	if (!TempFileOrErr)
680	return TempFileOrErr.takeError();
681	InputFile = *TempFileOrErr;
682	}
683
684	ObjcopyArgs.push_back(
685	Elt: SS.save(S: Twine("--add-section=") + OFFLOAD_BUNDLER_MAGIC_STR +
686	BundlerConfig.TargetNames[I] + "=" + InputFile));
687	ObjcopyArgs.push_back(
688	Elt: SS.save(S: Twine("--set-section-flags=") + OFFLOAD_BUNDLER_MAGIC_STR +
689	BundlerConfig.TargetNames[I] + "=readonly,exclude"));
690	}
691	ObjcopyArgs.push_back(Elt: "--");
692	ObjcopyArgs.push_back(
693	Elt: BundlerConfig.InputFileNames[BundlerConfig.HostInputIndex]);
694	ObjcopyArgs.push_back(Elt: BundlerConfig.OutputFileNames.front());
695
696	if (Error Err = executeObjcopy(Objcopy: BundlerConfig.ObjcopyPath, Args: ObjcopyArgs))
697	return Err;
698
699	return Error::success();
700	}
701
702	Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) final {
703	return Error::success();
704	}
705
706	private:
707	Error executeObjcopy(StringRef Objcopy, ArrayRef<StringRef> Args) {
708	// If the user asked for the commands to be printed out, we do that
709	// instead of executing it.
710	if (BundlerConfig.PrintExternalCommands) {
711	errs() << "\"" << Objcopy << "\"";
712	for (StringRef Arg : drop_begin(RangeOrContainer&: Args, N: 1))
713	errs() << " \"" << Arg << "\"";
714	errs() << "\n";
715	} else {
716	if (sys::ExecuteAndWait(Program: Objcopy, Args))
717	return createStringError(EC: inconvertibleErrorCode(),
718	S: "'llvm-objcopy' tool failed");
719	}
720	return Error::success();
721	}
722
723	Expected<std::string> getHostBundle(StringRef Input) {
724	TempFileHandlerRAII TempFiles;
725
726	auto ModifiedObjPathOrErr = TempFiles.Create(Contents: std::nullopt);
727	if (!ModifiedObjPathOrErr)
728	return ModifiedObjPathOrErr.takeError();
729	StringRef ModifiedObjPath = *ModifiedObjPathOrErr;
730
731	BumpPtrAllocator Alloc;
732	StringSaver SS{Alloc};
733	SmallVector<StringRef, 16> ObjcopyArgs{"llvm-objcopy"};
734
735	ObjcopyArgs.push_back(Elt: "--regex");
736	ObjcopyArgs.push_back(Elt: "--remove-section=__CLANG_OFFLOAD_BUNDLE__.*");
737	ObjcopyArgs.push_back(Elt: "--");
738
739	StringRef ObjcopyInputFileName;
740	// When unbundling an archive, the content of each object file in the
741	// archive is passed to this function by parameter Input, which is different
742	// from the content of the original input archive file, therefore it needs
743	// to be saved to a temporary file before passed to llvm-objcopy. Otherwise,
744	// Input is the same as the content of the original input file, therefore
745	// temporary file is not needed.
746	if (StringRef(BundlerConfig.FilesType).starts_with(Prefix: "a")) {
747	auto InputFileOrErr =
748	TempFiles.Create(Contents: ArrayRef<char>(Input.data(), Input.size()));
749	if (!InputFileOrErr)
750	return InputFileOrErr.takeError();
751	ObjcopyInputFileName = *InputFileOrErr;
752	} else
753	ObjcopyInputFileName = BundlerConfig.InputFileNames.front();
754
755	ObjcopyArgs.push_back(Elt: ObjcopyInputFileName);
756	ObjcopyArgs.push_back(Elt: ModifiedObjPath);
757
758	if (Error Err = executeObjcopy(Objcopy: BundlerConfig.ObjcopyPath, Args: ObjcopyArgs))
759	return std::move(Err);
760
761	auto BufOrErr = MemoryBuffer::getFile(Filename: ModifiedObjPath);
762	if (!BufOrErr)
763	return createStringError(EC: BufOrErr.getError(),
764	S: "Failed to read back the modified object file");
765
766	return BufOrErr->get()->getBuffer().str();
767	}
768	};
769
770	/// Handler for text files. The bundled file will have the following format.
771	///
772	/// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
773	/// Bundle 1
774	/// "Comment OFFLOAD_BUNDLER_MAGIC_STR__END__ triple"
775	/// ...
776	/// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
777	/// Bundle N
778	/// "Comment OFFLOAD_BUNDLER_MAGIC_STR__END__ triple"
779	class TextFileHandler final : public FileHandler {
780	/// String that begins a line comment.
781	StringRef Comment;
782
783	/// String that initiates a bundle.
784	std::string BundleStartString;
785
786	/// String that closes a bundle.
787	std::string BundleEndString;
788
789	/// Number of chars read from input.
790	size_t ReadChars = 0u;
791
792	protected:
793	Error ReadHeader(MemoryBuffer &Input) final { return Error::success(); }
794
795	Expected<std::optional<StringRef>>
796	ReadBundleStart(MemoryBuffer &Input) final {
797	StringRef FC = Input.getBuffer();
798
799	// Find start of the bundle.
800	ReadChars = FC.find(Str: BundleStartString, From: ReadChars);
801	if (ReadChars == FC.npos)
802	return std::nullopt;
803
804	// Get position of the triple.
805	size_t TripleStart = ReadChars = ReadChars + BundleStartString.size();
806
807	// Get position that closes the triple.
808	size_t TripleEnd = ReadChars = FC.find(Str: "\n", From: ReadChars);
809	if (TripleEnd == FC.npos)
810	return std::nullopt;
811
812	// Next time we read after the new line.
813	++ReadChars;
814
815	return StringRef(&FC.data()[TripleStart], TripleEnd - TripleStart);
816	}
817
818	Error ReadBundleEnd(MemoryBuffer &Input) final {
819	StringRef FC = Input.getBuffer();
820
821	// Read up to the next new line.
822	assert(FC[ReadChars] == '\n' && "The bundle should end with a new line.");
823
824	size_t TripleEnd = ReadChars = FC.find(Str: "\n", From: ReadChars + 1);
825	if (TripleEnd != FC.npos)
826	// Next time we read after the new line.
827	++ReadChars;
828
829	return Error::success();
830	}
831
832	Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
833	StringRef FC = Input.getBuffer();
834	size_t BundleStart = ReadChars;
835
836	// Find end of the bundle.
837	size_t BundleEnd = ReadChars = FC.find(Str: BundleEndString, From: ReadChars);
838
839	StringRef Bundle(&FC.data()[BundleStart], BundleEnd - BundleStart);
840	OS << Bundle;
841
842	return Error::success();
843	}
844
845	Error WriteHeader(raw_ostream &OS,
846	ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
847	return Error::success();
848	}
849
850	Error WriteBundleStart(raw_ostream &OS, StringRef TargetTriple) final {
851	OS << BundleStartString << TargetTriple << "\n";
852	return Error::success();
853	}
854
855	Error WriteBundleEnd(raw_ostream &OS, StringRef TargetTriple) final {
856	OS << BundleEndString << TargetTriple << "\n";
857	return Error::success();
858	}
859
860	Error WriteBundle(raw_ostream &OS, MemoryBuffer &Input) final {
861	OS << Input.getBuffer();
862	return Error::success();
863	}
864
865	public:
866	TextFileHandler(StringRef Comment) : Comment(Comment), ReadChars(0) {
867	BundleStartString =
868	"\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__START__ ";
869	BundleEndString =
870	"\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__END__ ";
871	}
872
873	Error listBundleIDsCallback(MemoryBuffer &Input,
874	const BundleInfo &Info) final {
875	// TODO: To list bundle IDs in a bundled text file we need to go through
876	// all bundles. The format of bundled text file may need to include a
877	// header if the performance of listing bundle IDs of bundled text file is
878	// important.
879	ReadChars = Input.getBuffer().find(Str: BundleEndString, From: ReadChars);
880	if (Error Err = ReadBundleEnd(Input))
881	return Err;
882	return Error::success();
883	}
884	};
885	} // namespace
886
887	/// Return an appropriate object file handler. We use the specific object
888	/// handler if we know how to deal with that format, otherwise we use a default
889	/// binary file handler.
890	static std::unique_ptr<FileHandler>
891	CreateObjectFileHandler(MemoryBuffer &FirstInput,
892	const OffloadBundlerConfig &BundlerConfig) {
893	// Check if the input file format is one that we know how to deal with.
894	Expected<std::unique_ptr<Binary>> BinaryOrErr = createBinary(Source: FirstInput);
895
896	// We only support regular object files. If failed to open the input as a
897	// known binary or this is not an object file use the default binary handler.
898	if (errorToBool(Err: BinaryOrErr.takeError()) || !isa<ObjectFile>(Val: *BinaryOrErr))
899	return std::make_unique<BinaryFileHandler>(args: BundlerConfig);
900
901	// Otherwise create an object file handler. The handler will be owned by the
902	// client of this function.
903	return std::make_unique<ObjectFileHandler>(
904	args: std::unique_ptr<ObjectFile>(cast<ObjectFile>(Val: BinaryOrErr->release())),
905	args: BundlerConfig);
906	}
907
908	/// Return an appropriate handler given the input files and options.
909	static Expected<std::unique_ptr<FileHandler>>
910	CreateFileHandler(MemoryBuffer &FirstInput,
911	const OffloadBundlerConfig &BundlerConfig) {
912	std::string FilesType = BundlerConfig.FilesType;
913
914	if (FilesType == "i")
915	return std::make_unique<TextFileHandler>(/*Comment=*/args: "//");
916	if (FilesType == "ii")
917	return std::make_unique<TextFileHandler>(/*Comment=*/args: "//");
918	if (FilesType == "cui")
919	return std::make_unique<TextFileHandler>(/*Comment=*/args: "//");
920	if (FilesType == "hipi")
921	return std::make_unique<TextFileHandler>(/*Comment=*/args: "//");
922	// TODO: `.d` should be eventually removed once `-M` and its variants are
923	// handled properly in offload compilation.
924	if (FilesType == "d")
925	return std::make_unique<TextFileHandler>(/*Comment=*/args: "#");
926	if (FilesType == "ll")
927	return std::make_unique<TextFileHandler>(/*Comment=*/args: ";");
928	if (FilesType == "bc")
929	return std::make_unique<BinaryFileHandler>(args: BundlerConfig);
930	if (FilesType == "s")
931	return std::make_unique<TextFileHandler>(/*Comment=*/args: "#");
932	if (FilesType == "o")
933	return CreateObjectFileHandler(FirstInput, BundlerConfig);
934	if (FilesType == "a")
935	return CreateObjectFileHandler(FirstInput, BundlerConfig);
936	if (FilesType == "gch")
937	return std::make_unique<BinaryFileHandler>(args: BundlerConfig);
938	if (FilesType == "ast")
939	return std::make_unique<BinaryFileHandler>(args: BundlerConfig);
940
941	return createStringError(EC: errc::invalid_argument,
942	S: "'" + FilesType + "': invalid file type specified");
943	}
944
945	OffloadBundlerConfig::OffloadBundlerConfig()
946	: CompressedBundleVersion(CompressedOffloadBundle::DefaultVersion) {
947	if (llvm::compression::zstd::isAvailable()) {
948	CompressionFormat = llvm::compression::Format::Zstd;
949	// Compression level 3 is usually sufficient for zstd since long distance
950	// matching is enabled.
951	CompressionLevel = 3;
952	} else if (llvm::compression::zlib::isAvailable()) {
953	CompressionFormat = llvm::compression::Format::Zlib;
954	// Use default level for zlib since higher level does not have significant
955	// improvement.
956	CompressionLevel = llvm::compression::zlib::DefaultCompression;
957	}
958	auto IgnoreEnvVarOpt =
959	llvm::sys::Process::GetEnv(name: "OFFLOAD_BUNDLER_IGNORE_ENV_VAR");
960	if (IgnoreEnvVarOpt.has_value() && IgnoreEnvVarOpt.value() == "1")
961	return;
962	auto VerboseEnvVarOpt = llvm::sys::Process::GetEnv(name: "OFFLOAD_BUNDLER_VERBOSE");
963	if (VerboseEnvVarOpt.has_value())
964	Verbose = VerboseEnvVarOpt.value() == "1";
965	auto CompressEnvVarOpt =
966	llvm::sys::Process::GetEnv(name: "OFFLOAD_BUNDLER_COMPRESS");
967	if (CompressEnvVarOpt.has_value())
968	Compress = CompressEnvVarOpt.value() == "1";
969	auto CompressionLevelEnvVarOpt =
970	llvm::sys::Process::GetEnv(name: "OFFLOAD_BUNDLER_COMPRESSION_LEVEL");
971	if (CompressionLevelEnvVarOpt.has_value()) {
972	llvm::StringRef CompressionLevelStr = CompressionLevelEnvVarOpt.value();
973	int Level;
974	if (!CompressionLevelStr.getAsInteger(Radix: 10, Result&: Level))
975	CompressionLevel = Level;
976	else
977	llvm::errs()
978	<< "Warning: Invalid value for OFFLOAD_BUNDLER_COMPRESSION_LEVEL: "
979	<< CompressionLevelStr.str() << ". Ignoring it.\n";
980	}
981	auto CompressedBundleFormatVersionOpt =
982	llvm::sys::Process::GetEnv(name: "COMPRESSED_BUNDLE_FORMAT_VERSION");
983	if (CompressedBundleFormatVersionOpt.has_value()) {
984	llvm::StringRef VersionStr = CompressedBundleFormatVersionOpt.value();
985	uint16_t Version;
986	if (!VersionStr.getAsInteger(Radix: 10, Result&: Version)) {
987	if (Version >= 2 && Version <= 3)
988	CompressedBundleVersion = Version;
989	else
990	llvm::errs()
991	<< "Warning: Invalid value for COMPRESSED_BUNDLE_FORMAT_VERSION: "
992	<< VersionStr.str()
993	<< ". Valid values are 2 or 3. Using default version "
994	<< CompressedBundleVersion << ".\n";
995	} else
996	llvm::errs()
997	<< "Warning: Invalid value for COMPRESSED_BUNDLE_FORMAT_VERSION: "
998	<< VersionStr.str() << ". Using default version "
999	<< CompressedBundleVersion << ".\n";
1000	}
1001	}
1002
1003	// Utility function to format numbers with commas
1004	static std::string formatWithCommas(unsigned long long Value) {
1005	std::string Num = std::to_string(val: Value);
1006	int InsertPosition = Num.length() - 3;
1007	while (InsertPosition > 0) {
1008	Num.insert(pos: InsertPosition, s: ",");
1009	InsertPosition -= 3;
1010	}
1011	return Num;
1012	}
1013
1014	llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
1015	CompressedOffloadBundle::compress(llvm::compression::Params P,
1016	const llvm::MemoryBuffer &Input,
1017	uint16_t Version, bool Verbose) {
1018	if (!llvm::compression::zstd::isAvailable() &&
1019	!llvm::compression::zlib::isAvailable())
1020	return createStringError(EC: llvm::inconvertibleErrorCode(),
1021	S: "Compression not supported");
1022	llvm::Timer HashTimer("Hash Calculation Timer", "Hash calculation time",
1023	*ClangOffloadBundlerTimerGroup);
1024	if (Verbose)
1025	HashTimer.startTimer();
1026	llvm::MD5 Hash;
1027	llvm::MD5::MD5Result Result;
1028	Hash.update(Str: Input.getBuffer());
1029	Hash.final(Result);
1030	uint64_t TruncatedHash = Result.low();
1031	if (Verbose)
1032	HashTimer.stopTimer();
1033
1034	SmallVector<uint8_t, 0> CompressedBuffer;
1035	auto BufferUint8 = llvm::ArrayRef<uint8_t>(
1036	reinterpret_cast<const uint8_t *>(Input.getBuffer().data()),
1037	Input.getBuffer().size());
1038	llvm::Timer CompressTimer("Compression Timer", "Compression time",
1039	*ClangOffloadBundlerTimerGroup);
1040	if (Verbose)
1041	CompressTimer.startTimer();
1042	llvm::compression::compress(P, Input: BufferUint8, Output&: CompressedBuffer);
1043	if (Verbose)
1044	CompressTimer.stopTimer();
1045
1046	uint16_t CompressionMethod = static_cast<uint16_t>(P.format);
1047
1048	// Store sizes in 64-bit variables first
1049	uint64_t UncompressedSize64 = Input.getBuffer().size();
1050	uint64_t TotalFileSize64;
1051
1052	// Calculate total file size based on version
1053	if (Version == 2) {
1054	// For V2, ensure the sizes don't exceed 32-bit limit
1055	if (UncompressedSize64 > std::numeric_limits<uint32_t>::max())
1056	return createStringError(EC: llvm::inconvertibleErrorCode(),
1057	S: "Uncompressed size exceeds version 2 limit");
1058	if ((MagicNumber.size() + sizeof(uint32_t) + sizeof(Version) +
1059	sizeof(CompressionMethod) + sizeof(uint32_t) + sizeof(TruncatedHash) +
1060	CompressedBuffer.size()) > std::numeric_limits<uint32_t>::max())
1061	return createStringError(EC: llvm::inconvertibleErrorCode(),
1062	S: "Total file size exceeds version 2 limit");
1063
1064	TotalFileSize64 = MagicNumber.size() + sizeof(uint32_t) + sizeof(Version) +
1065	sizeof(CompressionMethod) + sizeof(uint32_t) +
1066	sizeof(TruncatedHash) + CompressedBuffer.size();
1067	} else { // Version 3
1068	TotalFileSize64 = MagicNumber.size() + sizeof(uint64_t) + sizeof(Version) +
1069	sizeof(CompressionMethod) + sizeof(uint64_t) +
1070	sizeof(TruncatedHash) + CompressedBuffer.size();
1071	}
1072
1073	SmallVector<char, 0> FinalBuffer;
1074	llvm::raw_svector_ostream OS(FinalBuffer);
1075	OS << MagicNumber;
1076	OS.write(Ptr: reinterpret_cast<const char *>(&Version), Size: sizeof(Version));
1077	OS.write(Ptr: reinterpret_cast<const char *>(&CompressionMethod),
1078	Size: sizeof(CompressionMethod));
1079
1080	// Write size fields according to version
1081	if (Version == 2) {
1082	uint32_t TotalFileSize32 = static_cast<uint32_t>(TotalFileSize64);
1083	uint32_t UncompressedSize32 = static_cast<uint32_t>(UncompressedSize64);
1084	OS.write(Ptr: reinterpret_cast<const char *>(&TotalFileSize32),
1085	Size: sizeof(TotalFileSize32));
1086	OS.write(Ptr: reinterpret_cast<const char *>(&UncompressedSize32),
1087	Size: sizeof(UncompressedSize32));
1088	} else { // Version 3
1089	OS.write(Ptr: reinterpret_cast<const char *>(&TotalFileSize64),
1090	Size: sizeof(TotalFileSize64));
1091	OS.write(Ptr: reinterpret_cast<const char *>(&UncompressedSize64),
1092	Size: sizeof(UncompressedSize64));
1093	}
1094
1095	OS.write(Ptr: reinterpret_cast<const char *>(&TruncatedHash),
1096	Size: sizeof(TruncatedHash));
1097	OS.write(Ptr: reinterpret_cast<const char *>(CompressedBuffer.data()),
1098	Size: CompressedBuffer.size());
1099
1100	if (Verbose) {
1101	auto MethodUsed =
1102	P.format == llvm::compression::Format::Zstd ? "zstd" : "zlib";
1103	double CompressionRate =
1104	static_cast<double>(UncompressedSize64) / CompressedBuffer.size();
1105	double CompressionTimeSeconds = CompressTimer.getTotalTime().getWallTime();
1106	double CompressionSpeedMBs =
1107	(UncompressedSize64 / (1024.0 * 1024.0)) / CompressionTimeSeconds;
1108	llvm::errs() << "Compressed bundle format version: " << Version << "\n"
1109	<< "Total file size (including headers): "
1110	<< formatWithCommas(Value: TotalFileSize64) << " bytes\n"
1111	<< "Compression method used: " << MethodUsed << "\n"
1112	<< "Compression level: " << P.level << "\n"
1113	<< "Binary size before compression: "
1114	<< formatWithCommas(Value: UncompressedSize64) << " bytes\n"
1115	<< "Binary size after compression: "
1116	<< formatWithCommas(Value: CompressedBuffer.size()) << " bytes\n"
1117	<< "Compression rate: "
1118	<< llvm::format(Fmt: "%.2lf", Vals: CompressionRate) << "\n"
1119	<< "Compression ratio: "
1120	<< llvm::format(Fmt: "%.2lf%%", Vals: 100.0 / CompressionRate) << "\n"
1121	<< "Compression speed: "
1122	<< llvm::format(Fmt: "%.2lf MB/s", Vals: CompressionSpeedMBs) << "\n"
1123	<< "Truncated MD5 hash: "
1124	<< llvm::format_hex(N: TruncatedHash, Width: 16) << "\n";
1125	}
1126
1127	return llvm::MemoryBuffer::getMemBufferCopy(
1128	InputData: llvm::StringRef(FinalBuffer.data(), FinalBuffer.size()));
1129	}
1130
1131	// Use packed structs to avoid padding, such that the structs map the serialized
1132	// format.
1133	LLVM_PACKED_START
1134	union RawCompressedBundleHeader {
1135	struct CommonFields {
1136	uint32_t Magic;
1137	uint16_t Version;
1138	uint16_t Method;
1139	};
1140
1141	struct V1Header {
1142	CommonFields Common;
1143	uint32_t UncompressedFileSize;
1144	uint64_t Hash;
1145	};
1146
1147	struct V2Header {
1148	CommonFields Common;
1149	uint32_t FileSize;
1150	uint32_t UncompressedFileSize;
1151	uint64_t Hash;
1152	};
1153
1154	struct V3Header {
1155	CommonFields Common;
1156	uint64_t FileSize;
1157	uint64_t UncompressedFileSize;
1158	uint64_t Hash;
1159	};
1160
1161	CommonFields Common;
1162	V1Header V1;
1163	V2Header V2;
1164	V3Header V3;
1165	};
1166	LLVM_PACKED_END
1167
1168	// Helper method to get header size based on version
1169	static size_t getHeaderSize(uint16_t Version) {
1170	switch (Version) {
1171	case 1:
1172	return sizeof(RawCompressedBundleHeader::V1Header);
1173	case 2:
1174	return sizeof(RawCompressedBundleHeader::V2Header);
1175	case 3:
1176	return sizeof(RawCompressedBundleHeader::V3Header);
1177	default:
1178	llvm_unreachable("Unsupported version");
1179	}
1180	}
1181
1182	Expected<CompressedOffloadBundle::CompressedBundleHeader>
1183	CompressedOffloadBundle::CompressedBundleHeader::tryParse(StringRef Blob) {
1184	assert(Blob.size() >= sizeof(RawCompressedBundleHeader::CommonFields));
1185	assert(llvm::identify_magic(Blob) ==
1186	llvm::file_magic::offload_bundle_compressed);
1187
1188	RawCompressedBundleHeader Header;
1189	memcpy(dest: &Header, src: Blob.data(), n: std::min(a: Blob.size(), b: sizeof(Header)));
1190
1191	CompressedBundleHeader Normalized;
1192	Normalized.Version = Header.Common.Version;
1193
1194	size_t RequiredSize = getHeaderSize(Version: Normalized.Version);
1195	if (Blob.size() < RequiredSize)
1196	return createStringError(EC: inconvertibleErrorCode(),
1197	S: "Compressed bundle header size too small");
1198
1199	switch (Normalized.Version) {
1200	case 1:
1201	Normalized.UncompressedFileSize = Header.V1.UncompressedFileSize;
1202	Normalized.Hash = Header.V1.Hash;
1203	break;
1204	case 2:
1205	Normalized.FileSize = Header.V2.FileSize;
1206	Normalized.UncompressedFileSize = Header.V2.UncompressedFileSize;
1207	Normalized.Hash = Header.V2.Hash;
1208	break;
1209	case 3:
1210	Normalized.FileSize = Header.V3.FileSize;
1211	Normalized.UncompressedFileSize = Header.V3.UncompressedFileSize;
1212	Normalized.Hash = Header.V3.Hash;
1213	break;
1214	default:
1215	return createStringError(EC: inconvertibleErrorCode(),
1216	S: "Unknown compressed bundle version");
1217	}
1218
1219	// Determine compression format
1220	switch (Header.Common.Method) {
1221	case static_cast<uint16_t>(compression::Format::Zlib):
1222	case static_cast<uint16_t>(compression::Format::Zstd):
1223	Normalized.CompressionFormat =
1224	static_cast<compression::Format>(Header.Common.Method);
1225	break;
1226	default:
1227	return createStringError(EC: inconvertibleErrorCode(),
1228	S: "Unknown compressing method");
1229	}
1230
1231	return Normalized;
1232	}
1233
1234	llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
1235	CompressedOffloadBundle::decompress(const llvm::MemoryBuffer &Input,
1236	bool Verbose) {
1237	StringRef Blob = Input.getBuffer();
1238
1239	// Check minimum header size (using V1 as it's the smallest)
1240	if (Blob.size() < sizeof(RawCompressedBundleHeader::CommonFields))
1241	return llvm::MemoryBuffer::getMemBufferCopy(InputData: Blob);
1242
1243	if (llvm::identify_magic(magic: Blob) !=
1244	llvm::file_magic::offload_bundle_compressed) {
1245	if (Verbose)
1246	llvm::errs() << "Uncompressed bundle.\n";
1247	return llvm::MemoryBuffer::getMemBufferCopy(InputData: Blob);
1248	}
1249
1250	Expected<CompressedBundleHeader> HeaderOrErr =
1251	CompressedBundleHeader::tryParse(Blob);
1252	if (!HeaderOrErr)
1253	return HeaderOrErr.takeError();
1254
1255	const CompressedBundleHeader &Normalized = *HeaderOrErr;
1256	unsigned ThisVersion = Normalized.Version;
1257	size_t HeaderSize = getHeaderSize(Version: ThisVersion);
1258
1259	llvm::compression::Format CompressionFormat = Normalized.CompressionFormat;
1260
1261	size_t TotalFileSize = Normalized.FileSize.value_or(u: 0);
1262	size_t UncompressedSize = Normalized.UncompressedFileSize;
1263	auto StoredHash = Normalized.Hash;
1264
1265	llvm::Timer DecompressTimer("Decompression Timer", "Decompression time",
1266	*ClangOffloadBundlerTimerGroup);
1267	if (Verbose)
1268	DecompressTimer.startTimer();
1269
1270	SmallVector<uint8_t, 0> DecompressedData;
1271	StringRef CompressedData = Blob.substr(Start: HeaderSize);
1272	if (llvm::Error DecompressionError = llvm::compression::decompress(
1273	F: CompressionFormat, Input: llvm::arrayRefFromStringRef(Input: CompressedData),
1274	Output&: DecompressedData, UncompressedSize))
1275	return createStringError(EC: inconvertibleErrorCode(),
1276	S: "Could not decompress embedded file contents: " +
1277	llvm::toString(E: std::move(DecompressionError)));
1278
1279	if (Verbose) {
1280	DecompressTimer.stopTimer();
1281
1282	double DecompressionTimeSeconds =
1283	DecompressTimer.getTotalTime().getWallTime();
1284
1285	// Recalculate MD5 hash for integrity check
1286	llvm::Timer HashRecalcTimer("Hash Recalculation Timer",
1287	"Hash recalculation time",
1288	*ClangOffloadBundlerTimerGroup);
1289	HashRecalcTimer.startTimer();
1290	llvm::MD5 Hash;
1291	llvm::MD5::MD5Result Result;
1292	Hash.update(Data: llvm::ArrayRef<uint8_t>(DecompressedData.data(),
1293	DecompressedData.size()));
1294	Hash.final(Result);
1295	uint64_t RecalculatedHash = Result.low();
1296	HashRecalcTimer.stopTimer();
1297	bool HashMatch = (StoredHash == RecalculatedHash);
1298
1299	double CompressionRate =
1300	static_cast<double>(UncompressedSize) / CompressedData.size();
1301	double DecompressionSpeedMBs =
1302	(UncompressedSize / (1024.0 * 1024.0)) / DecompressionTimeSeconds;
1303
1304	llvm::errs() << "Compressed bundle format version: " << ThisVersion << "\n";
1305	if (ThisVersion >= 2)
1306	llvm::errs() << "Total file size (from header): "
1307	<< formatWithCommas(Value: TotalFileSize) << " bytes\n";
1308	llvm::errs() << "Decompression method: "
1309	<< (CompressionFormat == llvm::compression::Format::Zlib
1310	? "zlib"
1311	: "zstd")
1312	<< "\n"
1313	<< "Size before decompression: "
1314	<< formatWithCommas(Value: CompressedData.size()) << " bytes\n"
1315	<< "Size after decompression: "
1316	<< formatWithCommas(Value: UncompressedSize) << " bytes\n"
1317	<< "Compression rate: "
1318	<< llvm::format(Fmt: "%.2lf", Vals: CompressionRate) << "\n"
1319	<< "Compression ratio: "
1320	<< llvm::format(Fmt: "%.2lf%%", Vals: 100.0 / CompressionRate) << "\n"
1321	<< "Decompression speed: "
1322	<< llvm::format(Fmt: "%.2lf MB/s", Vals: DecompressionSpeedMBs) << "\n"
1323	<< "Stored hash: " << llvm::format_hex(N: StoredHash, Width: 16) << "\n"
1324	<< "Recalculated hash: "
1325	<< llvm::format_hex(N: RecalculatedHash, Width: 16) << "\n"
1326	<< "Hashes match: " << (HashMatch ? "Yes" : "No") << "\n";
1327	}
1328
1329	return llvm::MemoryBuffer::getMemBufferCopy(
1330	InputData: llvm::toStringRef(Input: DecompressedData));
1331	}
1332
1333	// List bundle IDs. Return true if an error was found.
1334	Error OffloadBundler::ListBundleIDsInFile(
1335	StringRef InputFileName, const OffloadBundlerConfig &BundlerConfig) {
1336	// Open Input file.
1337	ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
1338	MemoryBuffer::getFileOrSTDIN(Filename: InputFileName, /*IsText=*/true);
1339	if (std::error_code EC = CodeOrErr.getError())
1340	return createFileError(F: InputFileName, EC);
1341
1342	// Decompress the input if necessary.
1343	Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =
1344	CompressedOffloadBundle::decompress(Input: **CodeOrErr, Verbose: BundlerConfig.Verbose);
1345	if (!DecompressedBufferOrErr)
1346	return createStringError(
1347	EC: inconvertibleErrorCode(),
1348	S: "Failed to decompress input: " +
1349	llvm::toString(E: DecompressedBufferOrErr.takeError()));
1350
1351	MemoryBuffer &DecompressedInput = **DecompressedBufferOrErr;
1352
1353	// Select the right files handler.
1354	Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1355	CreateFileHandler(FirstInput&: DecompressedInput, BundlerConfig);
1356	if (!FileHandlerOrErr)
1357	return FileHandlerOrErr.takeError();
1358
1359	std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
1360	assert(FH);
1361	return FH->listBundleIDs(Input&: DecompressedInput);
1362	}
1363
1364	/// @brief Checks if a code object \p CodeObjectInfo is compatible with a given
1365	/// target \p TargetInfo.
1366	/// @link https://clang.llvm.org/docs/ClangOffloadBundler.html#bundle-entry-id
1367	bool isCodeObjectCompatible(const OffloadTargetInfo &CodeObjectInfo,
1368	const OffloadTargetInfo &TargetInfo) {
1369
1370	// Compatible in case of exact match.
1371	if (CodeObjectInfo == TargetInfo) {
1372	DEBUG_WITH_TYPE("CodeObjectCompatibility",
1373	dbgs() << "Compatible: Exact match: \t[CodeObject: "
1374	<< CodeObjectInfo.str()
1375	<< "]\t:\t[Target: " << TargetInfo.str() << "]\n");
1376	return true;
1377	}
1378
1379	// Incompatible if Kinds or Triples mismatch.
1380	if (!CodeObjectInfo.isOffloadKindCompatible(TargetOffloadKind: TargetInfo.OffloadKind) ||
1381	!CodeObjectInfo.Triple.isCompatibleWith(Other: TargetInfo.Triple)) {
1382	DEBUG_WITH_TYPE(
1383	"CodeObjectCompatibility",
1384	dbgs() << "Incompatible: Kind/Triple mismatch \t[CodeObject: "
1385	<< CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
1386	<< "]\n");
1387	return false;
1388	}
1389
1390	// Incompatible if Processors mismatch.
1391	llvm::StringMap<bool> CodeObjectFeatureMap, TargetFeatureMap;
1392	std::optional<StringRef> CodeObjectProc = clang::parseTargetID(
1393	T: CodeObjectInfo.Triple, OffloadArch: CodeObjectInfo.TargetID, FeatureMap: &CodeObjectFeatureMap);
1394	std::optional<StringRef> TargetProc = clang::parseTargetID(
1395	T: TargetInfo.Triple, OffloadArch: TargetInfo.TargetID, FeatureMap: &TargetFeatureMap);
1396
1397	// Both TargetProc and CodeObjectProc can't be empty here.
1398	if (!TargetProc || !CodeObjectProc ||
1399	CodeObjectProc.value() != TargetProc.value()) {
1400	DEBUG_WITH_TYPE("CodeObjectCompatibility",
1401	dbgs() << "Incompatible: Processor mismatch \t[CodeObject: "
1402	<< CodeObjectInfo.str()
1403	<< "]\t:\t[Target: " << TargetInfo.str() << "]\n");
1404	return false;
1405	}
1406
1407	// Incompatible if CodeObject has more features than Target, irrespective of
1408	// type or sign of features.
1409	if (CodeObjectFeatureMap.getNumItems() > TargetFeatureMap.getNumItems()) {
1410	DEBUG_WITH_TYPE("CodeObjectCompatibility",
1411	dbgs() << "Incompatible: CodeObject has more features "
1412	"than target \t[CodeObject: "
1413	<< CodeObjectInfo.str()
1414	<< "]\t:\t[Target: " << TargetInfo.str() << "]\n");
1415	return false;
1416	}
1417
1418	// Compatible if each target feature specified by target is compatible with
1419	// target feature of code object. The target feature is compatible if the
1420	// code object does not specify it (meaning Any), or if it specifies it
1421	// with the same value (meaning On or Off).
1422	for (const auto &CodeObjectFeature : CodeObjectFeatureMap) {
1423	auto TargetFeature = TargetFeatureMap.find(Key: CodeObjectFeature.getKey());
1424	if (TargetFeature == TargetFeatureMap.end()) {
1425	DEBUG_WITH_TYPE(
1426	"CodeObjectCompatibility",
1427	dbgs()
1428	<< "Incompatible: Value of CodeObject's non-ANY feature is "
1429	"not matching with Target feature's ANY value \t[CodeObject: "
1430	<< CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
1431	<< "]\n");
1432	return false;
1433	} else if (TargetFeature->getValue() != CodeObjectFeature.getValue()) {
1434	DEBUG_WITH_TYPE(
1435	"CodeObjectCompatibility",
1436	dbgs() << "Incompatible: Value of CodeObject's non-ANY feature is "
1437	"not matching with Target feature's non-ANY value "
1438	"\t[CodeObject: "
1439	<< CodeObjectInfo.str()
1440	<< "]\t:\t[Target: " << TargetInfo.str() << "]\n");
1441	return false;
1442	}
1443	}
1444
1445	// CodeObject is compatible if all features of Target are:
1446	// - either, present in the Code Object's features map with the same sign,
1447	// - or, the feature is missing from CodeObjects's features map i.e. it is
1448	// set to ANY
1449	DEBUG_WITH_TYPE(
1450	"CodeObjectCompatibility",
1451	dbgs() << "Compatible: Target IDs are compatible \t[CodeObject: "
1452	<< CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
1453	<< "]\n");
1454	return true;
1455	}
1456
1457	/// Bundle the files. Return true if an error was found.
1458	Error OffloadBundler::BundleFiles() {
1459	std::error_code EC;
1460
1461	// Create a buffer to hold the content before compressing.
1462	SmallVector<char, 0> Buffer;
1463	llvm::raw_svector_ostream BufferStream(Buffer);
1464
1465	// Open input files.
1466	SmallVector<std::unique_ptr<MemoryBuffer>, 8u> InputBuffers;
1467	InputBuffers.reserve(N: BundlerConfig.InputFileNames.size());
1468	for (auto &I : BundlerConfig.InputFileNames) {
1469	ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
1470	MemoryBuffer::getFileOrSTDIN(Filename: I, /*IsText=*/true);
1471	if (std::error_code EC = CodeOrErr.getError())
1472	return createFileError(F: I, EC);
1473	InputBuffers.emplace_back(Args: std::move(*CodeOrErr));
1474	}
1475
1476	// Get the file handler. We use the host buffer as reference.
1477	assert((BundlerConfig.HostInputIndex != ~0u || BundlerConfig.AllowNoHost) &&
1478	"Host input index undefined??");
1479	Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr = CreateFileHandler(
1480	FirstInput&: *InputBuffers[BundlerConfig.AllowNoHost ? 0
1481	: BundlerConfig.HostInputIndex],
1482	BundlerConfig);
1483	if (!FileHandlerOrErr)
1484	return FileHandlerOrErr.takeError();
1485
1486	std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
1487	assert(FH);
1488
1489	// Write header.
1490	if (Error Err = FH->WriteHeader(OS&: BufferStream, Inputs: InputBuffers))
1491	return Err;
1492
1493	// Write all bundles along with the start/end markers. If an error was found
1494	// writing the end of the bundle component, abort the bundle writing.
1495	auto Input = InputBuffers.begin();
1496	for (auto &Triple : BundlerConfig.TargetNames) {
1497	if (Error Err = FH->WriteBundleStart(OS&: BufferStream, TargetTriple: Triple))
1498	return Err;
1499	if (Error Err = FH->WriteBundle(OS&: BufferStream, Input&: **Input))
1500	return Err;
1501	if (Error Err = FH->WriteBundleEnd(OS&: BufferStream, TargetTriple: Triple))
1502	return Err;
1503	++Input;
1504	}
1505
1506	raw_fd_ostream OutputFile(BundlerConfig.OutputFileNames.front(), EC,
1507	sys::fs::OF_None);
1508	if (EC)
1509	return createFileError(F: BundlerConfig.OutputFileNames.front(), EC);
1510
1511	SmallVector<char, 0> CompressedBuffer;
1512	if (BundlerConfig.Compress) {
1513	std::unique_ptr<llvm::MemoryBuffer> BufferMemory =
1514	llvm::MemoryBuffer::getMemBufferCopy(
1515	InputData: llvm::StringRef(Buffer.data(), Buffer.size()));
1516	auto CompressionResult = CompressedOffloadBundle::compress(
1517	P: {BundlerConfig.CompressionFormat, BundlerConfig.CompressionLevel,
1518	/*zstdEnableLdm=*/true},
1519	Input: *BufferMemory, Version: BundlerConfig.CompressedBundleVersion,
1520	Verbose: BundlerConfig.Verbose);
1521	if (auto Error = CompressionResult.takeError())
1522	return Error;
1523
1524	auto CompressedMemBuffer = std::move(CompressionResult.get());
1525	CompressedBuffer.assign(in_start: CompressedMemBuffer->getBufferStart(),
1526	in_end: CompressedMemBuffer->getBufferEnd());
1527	} else
1528	CompressedBuffer = Buffer;
1529
1530	OutputFile.write(Ptr: CompressedBuffer.data(), Size: CompressedBuffer.size());
1531
1532	return FH->finalizeOutputFile();
1533	}
1534
1535	// Unbundle the files. Return true if an error was found.
1536	Error OffloadBundler::UnbundleFiles() {
1537	// Open Input file.
1538	ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
1539	MemoryBuffer::getFileOrSTDIN(Filename: BundlerConfig.InputFileNames.front(),
1540	/*IsText=*/true);
1541	if (std::error_code EC = CodeOrErr.getError())
1542	return createFileError(F: BundlerConfig.InputFileNames.front(), EC);
1543
1544	// Decompress the input if necessary.
1545	Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =
1546	CompressedOffloadBundle::decompress(Input: **CodeOrErr, Verbose: BundlerConfig.Verbose);
1547	if (!DecompressedBufferOrErr)
1548	return createStringError(
1549	EC: inconvertibleErrorCode(),
1550	S: "Failed to decompress input: " +
1551	llvm::toString(E: DecompressedBufferOrErr.takeError()));
1552
1553	MemoryBuffer &Input = **DecompressedBufferOrErr;
1554
1555	// Select the right files handler.
1556	Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1557	CreateFileHandler(FirstInput&: Input, BundlerConfig);
1558	if (!FileHandlerOrErr)
1559	return FileHandlerOrErr.takeError();
1560
1561	std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
1562	assert(FH);
1563
1564	// Read the header of the bundled file.
1565	if (Error Err = FH->ReadHeader(Input))
1566	return Err;
1567
1568	// Create a work list that consist of the map triple/output file.
1569	StringMap<StringRef> Worklist;
1570	auto Output = BundlerConfig.OutputFileNames.begin();
1571	for (auto &Triple : BundlerConfig.TargetNames) {
1572	if (!checkOffloadBundleID(Str: Triple))
1573	return createStringError(EC: errc::invalid_argument,
1574	S: "invalid bundle id from bundle config");
1575	Worklist[Triple] = *Output;
1576	++Output;
1577	}
1578
1579	// Read all the bundles that are in the work list. If we find no bundles we
1580	// assume the file is meant for the host target.
1581	bool FoundHostBundle = false;
1582	while (!Worklist.empty()) {
1583	Expected<std::optional<StringRef>> CurTripleOrErr =
1584	FH->ReadBundleStart(Input);
1585	if (!CurTripleOrErr)
1586	return CurTripleOrErr.takeError();
1587
1588	// We don't have more bundles.
1589	if (!*CurTripleOrErr)
1590	break;
1591
1592	StringRef CurTriple = **CurTripleOrErr;
1593	assert(!CurTriple.empty());
1594	if (!checkOffloadBundleID(Str: CurTriple))
1595	return createStringError(EC: errc::invalid_argument,
1596	S: "invalid bundle id read from the bundle");
1597
1598	auto Output = Worklist.begin();
1599	for (auto E = Worklist.end(); Output != E; Output++) {
1600	if (isCodeObjectCompatible(
1601	CodeObjectInfo: OffloadTargetInfo(CurTriple, BundlerConfig),
1602	TargetInfo: OffloadTargetInfo((*Output).first(), BundlerConfig))) {
1603	break;
1604	}
1605	}
1606
1607	if (Output == Worklist.end())
1608	continue;
1609	// Check if the output file can be opened and copy the bundle to it.
1610	std::error_code EC;
1611	raw_fd_ostream OutputFile((*Output).second, EC, sys::fs::OF_None);
1612	if (EC)
1613	return createFileError(F: (*Output).second, EC);
1614	if (Error Err = FH->ReadBundle(OS&: OutputFile, Input))
1615	return Err;
1616	if (Error Err = FH->ReadBundleEnd(Input))
1617	return Err;
1618	Worklist.erase(I: Output);
1619
1620	// Record if we found the host bundle.
1621	auto OffloadInfo = OffloadTargetInfo(CurTriple, BundlerConfig);
1622	if (OffloadInfo.hasHostKind())
1623	FoundHostBundle = true;
1624	}
1625
1626	if (!BundlerConfig.AllowMissingBundles && !Worklist.empty()) {
1627	std::string ErrMsg = "Can't find bundles for";
1628	std::set<StringRef> Sorted;
1629	for (auto &E : Worklist)
1630	Sorted.insert(x: E.first());
1631	unsigned I = 0;
1632	unsigned Last = Sorted.size() - 1;
1633	for (auto &E : Sorted) {
1634	if (I != 0 && Last > 1)
1635	ErrMsg += ",";
1636	ErrMsg += " ";
1637	if (I == Last && I != 0)
1638	ErrMsg += "and ";
1639	ErrMsg += E.str();
1640	++I;
1641	}
1642	return createStringError(EC: inconvertibleErrorCode(), S: ErrMsg);
1643	}
1644
1645	// If no bundles were found, assume the input file is the host bundle and
1646	// create empty files for the remaining targets.
1647	if (Worklist.size() == BundlerConfig.TargetNames.size()) {
1648	for (auto &E : Worklist) {
1649	std::error_code EC;
1650	raw_fd_ostream OutputFile(E.second, EC, sys::fs::OF_None);
1651	if (EC)
1652	return createFileError(F: E.second, EC);
1653
1654	// If this entry has a host kind, copy the input file to the output file.
1655	// We don't need to check E.getKey() here through checkOffloadBundleID
1656	// because the entire WorkList has been checked above.
1657	auto OffloadInfo = OffloadTargetInfo(E.getKey(), BundlerConfig);
1658	if (OffloadInfo.hasHostKind())
1659	OutputFile.write(Ptr: Input.getBufferStart(), Size: Input.getBufferSize());
1660	}
1661	return Error::success();
1662	}
1663
1664	// If we found elements, we emit an error if none of those were for the host
1665	// in case host bundle name was provided in command line.
1666	if (!(FoundHostBundle || BundlerConfig.HostInputIndex == ~0u ||
1667	BundlerConfig.AllowMissingBundles))
1668	return createStringError(EC: inconvertibleErrorCode(),
1669	S: "Can't find bundle for the host target");
1670
1671	// If we still have any elements in the worklist, create empty files for them.
1672	for (auto &E : Worklist) {
1673	std::error_code EC;
1674	raw_fd_ostream OutputFile(E.second, EC, sys::fs::OF_None);
1675	if (EC)
1676	return createFileError(F: E.second, EC);
1677	}
1678
1679	return Error::success();
1680	}
1681
1682	static Archive::Kind getDefaultArchiveKindForHost() {
1683	return Triple(sys::getDefaultTargetTriple()).isOSDarwin() ? Archive::K_DARWIN
1684	: Archive::K_GNU;
1685	}
1686
1687	/// @brief Computes a list of targets among all given targets which are
1688	/// compatible with this code object
1689	/// @param [in] CodeObjectInfo Code Object
1690	/// @param [out] CompatibleTargets List of all compatible targets among all
1691	/// given targets
1692	/// @return false, if no compatible target is found.
1693	static bool
1694	getCompatibleOffloadTargets(OffloadTargetInfo &CodeObjectInfo,
1695	SmallVectorImpl<StringRef> &CompatibleTargets,
1696	const OffloadBundlerConfig &BundlerConfig) {
1697	if (!CompatibleTargets.empty()) {
1698	DEBUG_WITH_TYPE("CodeObjectCompatibility",
1699	dbgs() << "CompatibleTargets list should be empty\n");
1700	return false;
1701	}
1702	for (auto &Target : BundlerConfig.TargetNames) {
1703	auto TargetInfo = OffloadTargetInfo(Target, BundlerConfig);
1704	if (isCodeObjectCompatible(CodeObjectInfo, TargetInfo))
1705	CompatibleTargets.push_back(Elt: Target);
1706	}
1707	return !CompatibleTargets.empty();
1708	}
1709
1710	// Check that each code object file in the input archive conforms to following
1711	// rule: for a specific processor, a feature either shows up in all target IDs,
1712	// or does not show up in any target IDs. Otherwise the target ID combination is
1713	// invalid.
1714	static Error
1715	CheckHeterogeneousArchive(StringRef ArchiveName,
1716	const OffloadBundlerConfig &BundlerConfig) {
1717	std::vector<std::unique_ptr<MemoryBuffer>> ArchiveBuffers;
1718	ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
1719	MemoryBuffer::getFileOrSTDIN(Filename: ArchiveName, IsText: true, RequiresNullTerminator: false);
1720	if (std::error_code EC = BufOrErr.getError())
1721	return createFileError(F: ArchiveName, EC);
1722
1723	ArchiveBuffers.push_back(x: std::move(*BufOrErr));
1724	Expected<std::unique_ptr<llvm::object::Archive>> LibOrErr =
1725	Archive::create(Source: ArchiveBuffers.back()->getMemBufferRef());
1726	if (!LibOrErr)
1727	return LibOrErr.takeError();
1728
1729	auto Archive = std::move(*LibOrErr);
1730
1731	Error ArchiveErr = Error::success();
1732	auto ChildEnd = Archive->child_end();
1733
1734	/// Iterate over all bundled code object files in the input archive.
1735	for (auto ArchiveIter = Archive->child_begin(Err&: ArchiveErr);
1736	ArchiveIter != ChildEnd; ++ArchiveIter) {
1737	if (ArchiveErr)
1738	return ArchiveErr;
1739	auto ArchiveChildNameOrErr = (*ArchiveIter).getName();
1740	if (!ArchiveChildNameOrErr)
1741	return ArchiveChildNameOrErr.takeError();
1742
1743	auto CodeObjectBufferRefOrErr = (*ArchiveIter).getMemoryBufferRef();
1744	if (!CodeObjectBufferRefOrErr)
1745	return CodeObjectBufferRefOrErr.takeError();
1746
1747	auto CodeObjectBuffer =
1748	MemoryBuffer::getMemBuffer(Ref: *CodeObjectBufferRefOrErr, RequiresNullTerminator: false);
1749
1750	Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1751	CreateFileHandler(FirstInput&: *CodeObjectBuffer, BundlerConfig);
1752	if (!FileHandlerOrErr)
1753	return FileHandlerOrErr.takeError();
1754
1755	std::unique_ptr<FileHandler> &FileHandler = *FileHandlerOrErr;
1756	assert(FileHandler);
1757
1758	std::set<StringRef> BundleIds;
1759	auto CodeObjectFileError =
1760	FileHandler->getBundleIDs(Input&: *CodeObjectBuffer, BundleIds);
1761	if (CodeObjectFileError)
1762	return CodeObjectFileError;
1763
1764	auto &&ConflictingArchs = clang::getConflictTargetIDCombination(TargetIDs: BundleIds);
1765	if (ConflictingArchs) {
1766	std::string ErrMsg =
1767	Twine("conflicting TargetIDs [" + ConflictingArchs.value().first +
1768	", " + ConflictingArchs.value().second + "] found in " +
1769	ArchiveChildNameOrErr.get() + " of " + ArchiveName)
1770	.str();
1771	return createStringError(EC: inconvertibleErrorCode(), S: ErrMsg);
1772	}
1773	}
1774
1775	return ArchiveErr;
1776	}
1777
1778	/// UnbundleArchive takes an archive file (".a") as input containing bundled
1779	/// code object files, and a list of offload targets (not host), and extracts
1780	/// the code objects into a new archive file for each offload target. Each
1781	/// resulting archive file contains all code object files corresponding to that
1782	/// particular offload target. The created archive file does not
1783	/// contain an index of the symbols and code object files are named as
1784	/// <<Parent Bundle Name>-<CodeObject's TargetID>>, with ':' replaced with '_'.
1785	Error OffloadBundler::UnbundleArchive() {
1786	std::vector<std::unique_ptr<MemoryBuffer>> ArchiveBuffers;
1787
1788	/// Map of target names with list of object files that will form the device
1789	/// specific archive for that target
1790	StringMap<std::vector<NewArchiveMember>> OutputArchivesMap;
1791
1792	// Map of target names and output archive filenames
1793	StringMap<StringRef> TargetOutputFileNameMap;
1794
1795	auto Output = BundlerConfig.OutputFileNames.begin();
1796	for (auto &Target : BundlerConfig.TargetNames) {
1797	TargetOutputFileNameMap[Target] = *Output;
1798	++Output;
1799	}
1800
1801	StringRef IFName = BundlerConfig.InputFileNames.front();
1802
1803	if (BundlerConfig.CheckInputArchive) {
1804	// For a specific processor, a feature either shows up in all target IDs, or
1805	// does not show up in any target IDs. Otherwise the target ID combination
1806	// is invalid.
1807	auto ArchiveError = CheckHeterogeneousArchive(ArchiveName: IFName, BundlerConfig);
1808	if (ArchiveError) {
1809	return ArchiveError;
1810	}
1811	}
1812
1813	ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
1814	MemoryBuffer::getFileOrSTDIN(Filename: IFName, IsText: true, RequiresNullTerminator: false);
1815	if (std::error_code EC = BufOrErr.getError())
1816	return createFileError(F: BundlerConfig.InputFileNames.front(), EC);
1817
1818	ArchiveBuffers.push_back(x: std::move(*BufOrErr));
1819	Expected<std::unique_ptr<llvm::object::Archive>> LibOrErr =
1820	Archive::create(Source: ArchiveBuffers.back()->getMemBufferRef());
1821	if (!LibOrErr)
1822	return LibOrErr.takeError();
1823
1824	auto Archive = std::move(*LibOrErr);
1825
1826	Error ArchiveErr = Error::success();
1827	auto ChildEnd = Archive->child_end();
1828
1829	/// Iterate over all bundled code object files in the input archive.
1830	for (auto ArchiveIter = Archive->child_begin(Err&: ArchiveErr);
1831	ArchiveIter != ChildEnd; ++ArchiveIter) {
1832	if (ArchiveErr)
1833	return ArchiveErr;
1834	auto ArchiveChildNameOrErr = (*ArchiveIter).getName();
1835	if (!ArchiveChildNameOrErr)
1836	return ArchiveChildNameOrErr.takeError();
1837
1838	StringRef BundledObjectFile = sys::path::filename(path: *ArchiveChildNameOrErr);
1839
1840	auto CodeObjectBufferRefOrErr = (*ArchiveIter).getMemoryBufferRef();
1841	if (!CodeObjectBufferRefOrErr)
1842	return CodeObjectBufferRefOrErr.takeError();
1843
1844	auto TempCodeObjectBuffer =
1845	MemoryBuffer::getMemBuffer(Ref: *CodeObjectBufferRefOrErr, RequiresNullTerminator: false);
1846
1847	// Decompress the buffer if necessary.
1848	Expected<std::unique_ptr<MemoryBuffer>> DecompressedBufferOrErr =
1849	CompressedOffloadBundle::decompress(Input: *TempCodeObjectBuffer,
1850	Verbose: BundlerConfig.Verbose);
1851	if (!DecompressedBufferOrErr)
1852	return createStringError(
1853	EC: inconvertibleErrorCode(),
1854	S: "Failed to decompress code object: " +
1855	llvm::toString(E: DecompressedBufferOrErr.takeError()));
1856
1857	MemoryBuffer &CodeObjectBuffer = **DecompressedBufferOrErr;
1858
1859	Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1860	CreateFileHandler(FirstInput&: CodeObjectBuffer, BundlerConfig);
1861	if (!FileHandlerOrErr)
1862	return FileHandlerOrErr.takeError();
1863
1864	std::unique_ptr<FileHandler> &FileHandler = *FileHandlerOrErr;
1865	assert(FileHandler &&
1866	"FileHandle creation failed for file in the archive!");
1867
1868	if (Error ReadErr = FileHandler->ReadHeader(Input&: CodeObjectBuffer))
1869	return ReadErr;
1870
1871	Expected<std::optional<StringRef>> CurBundleIDOrErr =
1872	FileHandler->ReadBundleStart(Input&: CodeObjectBuffer);
1873	if (!CurBundleIDOrErr)
1874	return CurBundleIDOrErr.takeError();
1875
1876	std::optional<StringRef> OptionalCurBundleID = *CurBundleIDOrErr;
1877	// No device code in this child, skip.
1878	if (!OptionalCurBundleID)
1879	continue;
1880	StringRef CodeObject = *OptionalCurBundleID;
1881
1882	// Process all bundle entries (CodeObjects) found in this child of input
1883	// archive.
1884	while (!CodeObject.empty()) {
1885	SmallVector<StringRef> CompatibleTargets;
1886	if (!checkOffloadBundleID(Str: CodeObject)) {
1887	return createStringError(EC: errc::invalid_argument,
1888	S: "Invalid bundle id read from code object");
1889	}
1890	auto CodeObjectInfo = OffloadTargetInfo(CodeObject, BundlerConfig);
1891	if (getCompatibleOffloadTargets(CodeObjectInfo, CompatibleTargets,
1892	BundlerConfig)) {
1893	std::string BundleData;
1894	raw_string_ostream DataStream(BundleData);
1895	if (Error Err = FileHandler->ReadBundle(OS&: DataStream, Input&: CodeObjectBuffer))
1896	return Err;
1897
1898	for (auto &CompatibleTarget : CompatibleTargets) {
1899	SmallString<128> BundledObjectFileName;
1900	BundledObjectFileName.assign(RHS: BundledObjectFile);
1901	auto OutputBundleName =
1902	Twine(llvm::sys::path::stem(path: BundledObjectFileName) + "-" +
1903	CodeObject +
1904	getDeviceLibraryFileName(BundleFileName: BundledObjectFileName,
1905	Device: CodeObjectInfo.TargetID))
1906	.str();
1907	// Replace ':' in optional target feature list with '_' to ensure
1908	// cross-platform validity.
1909	llvm::replace(Range&: OutputBundleName, OldValue: ':', NewValue: '_');
1910
1911	std::unique_ptr<MemoryBuffer> MemBuf = MemoryBuffer::getMemBufferCopy(
1912	InputData: DataStream.str(), BufferName: OutputBundleName);
1913	ArchiveBuffers.push_back(x: std::move(MemBuf));
1914	llvm::MemoryBufferRef MemBufRef =
1915	MemoryBufferRef(*(ArchiveBuffers.back()));
1916
1917	// For inserting <CompatibleTarget, list<CodeObject>> entry in
1918	// OutputArchivesMap.
1919	OutputArchivesMap[CompatibleTarget].push_back(
1920	x: NewArchiveMember(MemBufRef));
1921	}
1922	}
1923
1924	if (Error Err = FileHandler->ReadBundleEnd(Input&: CodeObjectBuffer))
1925	return Err;
1926
1927	Expected<std::optional<StringRef>> NextTripleOrErr =
1928	FileHandler->ReadBundleStart(Input&: CodeObjectBuffer);
1929	if (!NextTripleOrErr)
1930	return NextTripleOrErr.takeError();
1931
1932	CodeObject = ((*NextTripleOrErr).has_value()) ? **NextTripleOrErr : "";
1933	} // End of processing of all bundle entries of this child of input archive.
1934	} // End of while over children of input archive.
1935
1936	assert(!ArchiveErr && "Error occurred while reading archive!");
1937
1938	/// Write out an archive for each target
1939	for (auto &Target : BundlerConfig.TargetNames) {
1940	StringRef FileName = TargetOutputFileNameMap[Target];
1941	StringMapIterator<std::vector<llvm::NewArchiveMember>> CurArchiveMembers =
1942	OutputArchivesMap.find(Key: Target);
1943	if (CurArchiveMembers != OutputArchivesMap.end()) {
1944	if (Error WriteErr = writeArchive(ArcName: FileName, NewMembers: CurArchiveMembers->getValue(),
1945	WriteSymtab: SymtabWritingMode::NormalSymtab,
1946	Kind: getDefaultArchiveKindForHost(), Deterministic: true,
1947	Thin: false, OldArchiveBuf: nullptr))
1948	return WriteErr;
1949	} else if (!BundlerConfig.AllowMissingBundles) {
1950	std::string ErrMsg =
1951	Twine("no compatible code object found for the target '" + Target +
1952	"' in heterogeneous archive library: " + IFName)
1953	.str();
1954	return createStringError(EC: inconvertibleErrorCode(), S: ErrMsg);
1955	} else { // Create an empty archive file if no compatible code object is
1956	// found and "allow-missing-bundles" is enabled. It ensures that
1957	// the linker using output of this step doesn't complain about
1958	// the missing input file.
1959	std::vector<llvm::NewArchiveMember> EmptyArchive;
1960	EmptyArchive.clear();
1961	if (Error WriteErr = writeArchive(
1962	ArcName: FileName, NewMembers: EmptyArchive, WriteSymtab: SymtabWritingMode::NormalSymtab,
1963	Kind: getDefaultArchiveKindForHost(), Deterministic: true, Thin: false, OldArchiveBuf: nullptr))
1964	return WriteErr;
1965	}
1966	}
1967
1968	return Error::success();
1969	}
1970
1971	bool clang::checkOffloadBundleID(const llvm::StringRef Str) {
1972	// <kind>-<triple>[-<target id>[:target features]]
1973	// <triple> := <arch>-<vendor>-<os>-<env>
1974	SmallVector<StringRef, 6> Components;
1975	Str.split(A&: Components, Separator: '-', /*MaxSplit=*/5);
1976	return Components.size() == 5 || Components.size() == 6;
1977	}
1978