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