xref: /freebsd-src/contrib/llvm-project/clang/lib/Driver/OffloadBundler.cpp (revision 06c3fb2749bda94cb5201f81ffdb8fa6c3161b2e)
1bdd1243dSDimitry Andric //===- OffloadBundler.cpp - File Bundling and Unbundling ------------------===//
2bdd1243dSDimitry Andric //
3bdd1243dSDimitry Andric // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4bdd1243dSDimitry Andric // See https://llvm.org/LICENSE.txt for license information.
5bdd1243dSDimitry Andric // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6bdd1243dSDimitry Andric //
7bdd1243dSDimitry Andric //===----------------------------------------------------------------------===//
8bdd1243dSDimitry Andric ///
9bdd1243dSDimitry Andric /// \file
10bdd1243dSDimitry Andric /// This file implements an offload bundling API that bundles different files
11bdd1243dSDimitry Andric /// that relate with the same source code but different targets into a single
12bdd1243dSDimitry Andric /// one. Also the implements the opposite functionality, i.e. unbundle files
13bdd1243dSDimitry Andric /// previous created by this API.
14bdd1243dSDimitry Andric ///
15bdd1243dSDimitry Andric //===----------------------------------------------------------------------===//
16bdd1243dSDimitry Andric 
17bdd1243dSDimitry Andric #include "clang/Driver/OffloadBundler.h"
18bdd1243dSDimitry Andric #include "clang/Basic/Cuda.h"
19bdd1243dSDimitry Andric #include "clang/Basic/TargetID.h"
20bdd1243dSDimitry Andric #include "clang/Basic/Version.h"
21bdd1243dSDimitry Andric #include "llvm/ADT/ArrayRef.h"
22bdd1243dSDimitry Andric #include "llvm/ADT/SmallString.h"
23bdd1243dSDimitry Andric #include "llvm/ADT/SmallVector.h"
24bdd1243dSDimitry Andric #include "llvm/ADT/StringMap.h"
25bdd1243dSDimitry Andric #include "llvm/ADT/StringRef.h"
26bdd1243dSDimitry Andric #include "llvm/Object/Archive.h"
27bdd1243dSDimitry Andric #include "llvm/Object/ArchiveWriter.h"
28bdd1243dSDimitry Andric #include "llvm/Object/Binary.h"
29bdd1243dSDimitry Andric #include "llvm/Object/ObjectFile.h"
30bdd1243dSDimitry Andric #include "llvm/Support/Casting.h"
31bdd1243dSDimitry Andric #include "llvm/Support/Debug.h"
32bdd1243dSDimitry Andric #include "llvm/Support/EndianStream.h"
33bdd1243dSDimitry Andric #include "llvm/Support/Errc.h"
34bdd1243dSDimitry Andric #include "llvm/Support/Error.h"
35bdd1243dSDimitry Andric #include "llvm/Support/ErrorOr.h"
36bdd1243dSDimitry Andric #include "llvm/Support/FileSystem.h"
37bdd1243dSDimitry Andric #include "llvm/Support/MemoryBuffer.h"
38bdd1243dSDimitry Andric #include "llvm/Support/Path.h"
39bdd1243dSDimitry Andric #include "llvm/Support/Program.h"
40bdd1243dSDimitry Andric #include "llvm/Support/Signals.h"
41bdd1243dSDimitry Andric #include "llvm/Support/StringSaver.h"
42bdd1243dSDimitry Andric #include "llvm/Support/WithColor.h"
43bdd1243dSDimitry Andric #include "llvm/Support/raw_ostream.h"
44*06c3fb27SDimitry Andric #include "llvm/TargetParser/Host.h"
45*06c3fb27SDimitry Andric #include "llvm/TargetParser/Triple.h"
46bdd1243dSDimitry Andric #include <algorithm>
47bdd1243dSDimitry Andric #include <cassert>
48bdd1243dSDimitry Andric #include <cstddef>
49bdd1243dSDimitry Andric #include <cstdint>
50bdd1243dSDimitry Andric #include <forward_list>
51bdd1243dSDimitry Andric #include <memory>
52bdd1243dSDimitry Andric #include <set>
53bdd1243dSDimitry Andric #include <string>
54bdd1243dSDimitry Andric #include <system_error>
55bdd1243dSDimitry Andric #include <utility>
56bdd1243dSDimitry Andric 
57bdd1243dSDimitry Andric using namespace llvm;
58bdd1243dSDimitry Andric using namespace llvm::object;
59bdd1243dSDimitry Andric using namespace clang;
60bdd1243dSDimitry Andric 
61bdd1243dSDimitry Andric /// Magic string that marks the existence of offloading data.
62bdd1243dSDimitry Andric #define OFFLOAD_BUNDLER_MAGIC_STR "__CLANG_OFFLOAD_BUNDLE__"
63bdd1243dSDimitry Andric 
64bdd1243dSDimitry Andric OffloadTargetInfo::OffloadTargetInfo(const StringRef Target,
65bdd1243dSDimitry Andric                                      const OffloadBundlerConfig &BC)
66bdd1243dSDimitry Andric     : BundlerConfig(BC) {
67bdd1243dSDimitry Andric 
68bdd1243dSDimitry Andric   // TODO: Add error checking from ClangOffloadBundler.cpp
69bdd1243dSDimitry Andric   auto TargetFeatures = Target.split(':');
70bdd1243dSDimitry Andric   auto TripleOrGPU = TargetFeatures.first.rsplit('-');
71bdd1243dSDimitry Andric 
72bdd1243dSDimitry Andric   if (clang::StringToCudaArch(TripleOrGPU.second) != clang::CudaArch::UNKNOWN) {
73bdd1243dSDimitry Andric     auto KindTriple = TripleOrGPU.first.split('-');
74bdd1243dSDimitry Andric     this->OffloadKind = KindTriple.first;
75*06c3fb27SDimitry Andric 
76*06c3fb27SDimitry Andric     // Enforce optional env field to standardize bundles
77*06c3fb27SDimitry Andric     llvm::Triple t = llvm::Triple(KindTriple.second);
78*06c3fb27SDimitry Andric     this->Triple = llvm::Triple(t.getArchName(), t.getVendorName(),
79*06c3fb27SDimitry Andric                                 t.getOSName(), t.getEnvironmentName());
80*06c3fb27SDimitry Andric 
81bdd1243dSDimitry Andric     this->TargetID = Target.substr(Target.find(TripleOrGPU.second));
82bdd1243dSDimitry Andric   } else {
83bdd1243dSDimitry Andric     auto KindTriple = TargetFeatures.first.split('-');
84bdd1243dSDimitry Andric     this->OffloadKind = KindTriple.first;
85*06c3fb27SDimitry Andric 
86*06c3fb27SDimitry Andric     // Enforce optional env field to standardize bundles
87*06c3fb27SDimitry Andric     llvm::Triple t = llvm::Triple(KindTriple.second);
88*06c3fb27SDimitry Andric     this->Triple = llvm::Triple(t.getArchName(), t.getVendorName(),
89*06c3fb27SDimitry Andric                                 t.getOSName(), t.getEnvironmentName());
90*06c3fb27SDimitry Andric 
91bdd1243dSDimitry Andric     this->TargetID = "";
92bdd1243dSDimitry Andric   }
93bdd1243dSDimitry Andric }
94bdd1243dSDimitry Andric 
95bdd1243dSDimitry Andric bool OffloadTargetInfo::hasHostKind() const {
96bdd1243dSDimitry Andric   return this->OffloadKind == "host";
97bdd1243dSDimitry Andric }
98bdd1243dSDimitry Andric 
99bdd1243dSDimitry Andric bool OffloadTargetInfo::isOffloadKindValid() const {
100bdd1243dSDimitry Andric   return OffloadKind == "host" || OffloadKind == "openmp" ||
101bdd1243dSDimitry Andric          OffloadKind == "hip" || OffloadKind == "hipv4";
102bdd1243dSDimitry Andric }
103bdd1243dSDimitry Andric 
104bdd1243dSDimitry Andric bool OffloadTargetInfo::isOffloadKindCompatible(
105bdd1243dSDimitry Andric     const StringRef TargetOffloadKind) const {
106bdd1243dSDimitry Andric   if (OffloadKind == TargetOffloadKind)
107bdd1243dSDimitry Andric     return true;
108bdd1243dSDimitry Andric   if (BundlerConfig.HipOpenmpCompatible) {
109*06c3fb27SDimitry Andric     bool HIPCompatibleWithOpenMP = OffloadKind.starts_with_insensitive("hip") &&
110bdd1243dSDimitry Andric                                    TargetOffloadKind == "openmp";
111bdd1243dSDimitry Andric     bool OpenMPCompatibleWithHIP =
112bdd1243dSDimitry Andric         OffloadKind == "openmp" &&
113*06c3fb27SDimitry Andric         TargetOffloadKind.starts_with_insensitive("hip");
114bdd1243dSDimitry Andric     return HIPCompatibleWithOpenMP || OpenMPCompatibleWithHIP;
115bdd1243dSDimitry Andric   }
116bdd1243dSDimitry Andric   return false;
117bdd1243dSDimitry Andric }
118bdd1243dSDimitry Andric 
119bdd1243dSDimitry Andric bool OffloadTargetInfo::isTripleValid() const {
120bdd1243dSDimitry Andric   return !Triple.str().empty() && Triple.getArch() != Triple::UnknownArch;
121bdd1243dSDimitry Andric }
122bdd1243dSDimitry Andric 
123bdd1243dSDimitry Andric bool OffloadTargetInfo::operator==(const OffloadTargetInfo &Target) const {
124bdd1243dSDimitry Andric   return OffloadKind == Target.OffloadKind &&
125bdd1243dSDimitry Andric          Triple.isCompatibleWith(Target.Triple) && TargetID == Target.TargetID;
126bdd1243dSDimitry Andric }
127bdd1243dSDimitry Andric 
128bdd1243dSDimitry Andric std::string OffloadTargetInfo::str() const {
129bdd1243dSDimitry Andric   return Twine(OffloadKind + "-" + Triple.str() + "-" + TargetID).str();
130bdd1243dSDimitry Andric }
131bdd1243dSDimitry Andric 
132bdd1243dSDimitry Andric static StringRef getDeviceFileExtension(StringRef Device,
133bdd1243dSDimitry Andric                                         StringRef BundleFileName) {
134bdd1243dSDimitry Andric   if (Device.contains("gfx"))
135bdd1243dSDimitry Andric     return ".bc";
136bdd1243dSDimitry Andric   if (Device.contains("sm_"))
137bdd1243dSDimitry Andric     return ".cubin";
138bdd1243dSDimitry Andric   return sys::path::extension(BundleFileName);
139bdd1243dSDimitry Andric }
140bdd1243dSDimitry Andric 
141bdd1243dSDimitry Andric static std::string getDeviceLibraryFileName(StringRef BundleFileName,
142bdd1243dSDimitry Andric                                             StringRef Device) {
143bdd1243dSDimitry Andric   StringRef LibName = sys::path::stem(BundleFileName);
144bdd1243dSDimitry Andric   StringRef Extension = getDeviceFileExtension(Device, BundleFileName);
145bdd1243dSDimitry Andric 
146bdd1243dSDimitry Andric   std::string Result;
147bdd1243dSDimitry Andric   Result += LibName;
148bdd1243dSDimitry Andric   Result += Extension;
149bdd1243dSDimitry Andric   return Result;
150bdd1243dSDimitry Andric }
151bdd1243dSDimitry Andric 
152bdd1243dSDimitry Andric /// @brief Checks if a code object \p CodeObjectInfo is compatible with a given
153bdd1243dSDimitry Andric /// target \p TargetInfo.
154bdd1243dSDimitry Andric /// @link https://clang.llvm.org/docs/ClangOffloadBundler.html#bundle-entry-id
155bdd1243dSDimitry Andric bool isCodeObjectCompatible(const OffloadTargetInfo &CodeObjectInfo,
156bdd1243dSDimitry Andric                             const OffloadTargetInfo &TargetInfo) {
157bdd1243dSDimitry Andric 
158bdd1243dSDimitry Andric   // Compatible in case of exact match.
159bdd1243dSDimitry Andric   if (CodeObjectInfo == TargetInfo) {
160bdd1243dSDimitry Andric     DEBUG_WITH_TYPE("CodeObjectCompatibility",
161bdd1243dSDimitry Andric                     dbgs() << "Compatible: Exact match: \t[CodeObject: "
162bdd1243dSDimitry Andric                            << CodeObjectInfo.str()
163bdd1243dSDimitry Andric                            << "]\t:\t[Target: " << TargetInfo.str() << "]\n");
164bdd1243dSDimitry Andric     return true;
165bdd1243dSDimitry Andric   }
166bdd1243dSDimitry Andric 
167bdd1243dSDimitry Andric   // Incompatible if Kinds or Triples mismatch.
168bdd1243dSDimitry Andric   if (!CodeObjectInfo.isOffloadKindCompatible(TargetInfo.OffloadKind) ||
169bdd1243dSDimitry Andric       !CodeObjectInfo.Triple.isCompatibleWith(TargetInfo.Triple)) {
170bdd1243dSDimitry Andric     DEBUG_WITH_TYPE(
171bdd1243dSDimitry Andric         "CodeObjectCompatibility",
172bdd1243dSDimitry Andric         dbgs() << "Incompatible: Kind/Triple mismatch \t[CodeObject: "
173bdd1243dSDimitry Andric                << CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
174bdd1243dSDimitry Andric                << "]\n");
175bdd1243dSDimitry Andric     return false;
176bdd1243dSDimitry Andric   }
177bdd1243dSDimitry Andric 
178bdd1243dSDimitry Andric   // Incompatible if target IDs are incompatible.
179bdd1243dSDimitry Andric   if (!clang::isCompatibleTargetID(CodeObjectInfo.TargetID,
180bdd1243dSDimitry Andric                                    TargetInfo.TargetID)) {
181bdd1243dSDimitry Andric     DEBUG_WITH_TYPE(
182bdd1243dSDimitry Andric         "CodeObjectCompatibility",
183bdd1243dSDimitry Andric         dbgs() << "Incompatible: target IDs are incompatible \t[CodeObject: "
184bdd1243dSDimitry Andric                << CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
185bdd1243dSDimitry Andric                << "]\n");
186bdd1243dSDimitry Andric     return false;
187bdd1243dSDimitry Andric   }
188bdd1243dSDimitry Andric 
189bdd1243dSDimitry Andric   DEBUG_WITH_TYPE(
190bdd1243dSDimitry Andric       "CodeObjectCompatibility",
191bdd1243dSDimitry Andric       dbgs() << "Compatible: Code Objects are compatible \t[CodeObject: "
192bdd1243dSDimitry Andric              << CodeObjectInfo.str() << "]\t:\t[Target: " << TargetInfo.str()
193bdd1243dSDimitry Andric              << "]\n");
194bdd1243dSDimitry Andric   return true;
195bdd1243dSDimitry Andric }
196bdd1243dSDimitry Andric 
197bdd1243dSDimitry Andric namespace {
198bdd1243dSDimitry Andric /// Generic file handler interface.
199bdd1243dSDimitry Andric class FileHandler {
200bdd1243dSDimitry Andric public:
201bdd1243dSDimitry Andric   struct BundleInfo {
202bdd1243dSDimitry Andric     StringRef BundleID;
203bdd1243dSDimitry Andric   };
204bdd1243dSDimitry Andric 
205bdd1243dSDimitry Andric   FileHandler() {}
206bdd1243dSDimitry Andric 
207bdd1243dSDimitry Andric   virtual ~FileHandler() {}
208bdd1243dSDimitry Andric 
209bdd1243dSDimitry Andric   /// Update the file handler with information from the header of the bundled
210bdd1243dSDimitry Andric   /// file.
211bdd1243dSDimitry Andric   virtual Error ReadHeader(MemoryBuffer &Input) = 0;
212bdd1243dSDimitry Andric 
213bdd1243dSDimitry Andric   /// Read the marker of the next bundled to be read in the file. The bundle
214bdd1243dSDimitry Andric   /// name is returned if there is one in the file, or `std::nullopt` if there
215bdd1243dSDimitry Andric   /// are no more bundles to be read.
216bdd1243dSDimitry Andric   virtual Expected<std::optional<StringRef>>
217bdd1243dSDimitry Andric   ReadBundleStart(MemoryBuffer &Input) = 0;
218bdd1243dSDimitry Andric 
219bdd1243dSDimitry Andric   /// Read the marker that closes the current bundle.
220bdd1243dSDimitry Andric   virtual Error ReadBundleEnd(MemoryBuffer &Input) = 0;
221bdd1243dSDimitry Andric 
222bdd1243dSDimitry Andric   /// Read the current bundle and write the result into the stream \a OS.
223bdd1243dSDimitry Andric   virtual Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) = 0;
224bdd1243dSDimitry Andric 
225bdd1243dSDimitry Andric   /// Write the header of the bundled file to \a OS based on the information
226bdd1243dSDimitry Andric   /// gathered from \a Inputs.
227bdd1243dSDimitry Andric   virtual Error WriteHeader(raw_fd_ostream &OS,
228bdd1243dSDimitry Andric                             ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) = 0;
229bdd1243dSDimitry Andric 
230bdd1243dSDimitry Andric   /// Write the marker that initiates a bundle for the triple \a TargetTriple to
231bdd1243dSDimitry Andric   /// \a OS.
232bdd1243dSDimitry Andric   virtual Error WriteBundleStart(raw_fd_ostream &OS,
233bdd1243dSDimitry Andric                                  StringRef TargetTriple) = 0;
234bdd1243dSDimitry Andric 
235bdd1243dSDimitry Andric   /// Write the marker that closes a bundle for the triple \a TargetTriple to \a
236bdd1243dSDimitry Andric   /// OS.
237bdd1243dSDimitry Andric   virtual Error WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) = 0;
238bdd1243dSDimitry Andric 
239bdd1243dSDimitry Andric   /// Write the bundle from \a Input into \a OS.
240bdd1243dSDimitry Andric   virtual Error WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) = 0;
241bdd1243dSDimitry Andric 
242bdd1243dSDimitry Andric   /// List bundle IDs in \a Input.
243bdd1243dSDimitry Andric   virtual Error listBundleIDs(MemoryBuffer &Input) {
244bdd1243dSDimitry Andric     if (Error Err = ReadHeader(Input))
245bdd1243dSDimitry Andric       return Err;
246bdd1243dSDimitry Andric     return forEachBundle(Input, [&](const BundleInfo &Info) -> Error {
247bdd1243dSDimitry Andric       llvm::outs() << Info.BundleID << '\n';
248bdd1243dSDimitry Andric       Error Err = listBundleIDsCallback(Input, Info);
249bdd1243dSDimitry Andric       if (Err)
250bdd1243dSDimitry Andric         return Err;
251bdd1243dSDimitry Andric       return Error::success();
252bdd1243dSDimitry Andric     });
253bdd1243dSDimitry Andric   }
254bdd1243dSDimitry Andric 
255bdd1243dSDimitry Andric   /// For each bundle in \a Input, do \a Func.
256bdd1243dSDimitry Andric   Error forEachBundle(MemoryBuffer &Input,
257bdd1243dSDimitry Andric                       std::function<Error(const BundleInfo &)> Func) {
258bdd1243dSDimitry Andric     while (true) {
259bdd1243dSDimitry Andric       Expected<std::optional<StringRef>> CurTripleOrErr =
260bdd1243dSDimitry Andric           ReadBundleStart(Input);
261bdd1243dSDimitry Andric       if (!CurTripleOrErr)
262bdd1243dSDimitry Andric         return CurTripleOrErr.takeError();
263bdd1243dSDimitry Andric 
264bdd1243dSDimitry Andric       // No more bundles.
265bdd1243dSDimitry Andric       if (!*CurTripleOrErr)
266bdd1243dSDimitry Andric         break;
267bdd1243dSDimitry Andric 
268bdd1243dSDimitry Andric       StringRef CurTriple = **CurTripleOrErr;
269bdd1243dSDimitry Andric       assert(!CurTriple.empty());
270bdd1243dSDimitry Andric 
271bdd1243dSDimitry Andric       BundleInfo Info{CurTriple};
272bdd1243dSDimitry Andric       if (Error Err = Func(Info))
273bdd1243dSDimitry Andric         return Err;
274bdd1243dSDimitry Andric     }
275bdd1243dSDimitry Andric     return Error::success();
276bdd1243dSDimitry Andric   }
277bdd1243dSDimitry Andric 
278bdd1243dSDimitry Andric protected:
279bdd1243dSDimitry Andric   virtual Error listBundleIDsCallback(MemoryBuffer &Input,
280bdd1243dSDimitry Andric                                       const BundleInfo &Info) {
281bdd1243dSDimitry Andric     return Error::success();
282bdd1243dSDimitry Andric   }
283bdd1243dSDimitry Andric };
284bdd1243dSDimitry Andric 
285bdd1243dSDimitry Andric /// Handler for binary files. The bundled file will have the following format
286bdd1243dSDimitry Andric /// (all integers are stored in little-endian format):
287bdd1243dSDimitry Andric ///
288bdd1243dSDimitry Andric /// "OFFLOAD_BUNDLER_MAGIC_STR" (ASCII encoding of the string)
289bdd1243dSDimitry Andric ///
290bdd1243dSDimitry Andric /// NumberOfOffloadBundles (8-byte integer)
291bdd1243dSDimitry Andric ///
292bdd1243dSDimitry Andric /// OffsetOfBundle1 (8-byte integer)
293bdd1243dSDimitry Andric /// SizeOfBundle1 (8-byte integer)
294bdd1243dSDimitry Andric /// NumberOfBytesInTripleOfBundle1 (8-byte integer)
295bdd1243dSDimitry Andric /// TripleOfBundle1 (byte length defined before)
296bdd1243dSDimitry Andric ///
297bdd1243dSDimitry Andric /// ...
298bdd1243dSDimitry Andric ///
299bdd1243dSDimitry Andric /// OffsetOfBundleN (8-byte integer)
300bdd1243dSDimitry Andric /// SizeOfBundleN (8-byte integer)
301bdd1243dSDimitry Andric /// NumberOfBytesInTripleOfBundleN (8-byte integer)
302bdd1243dSDimitry Andric /// TripleOfBundleN (byte length defined before)
303bdd1243dSDimitry Andric ///
304bdd1243dSDimitry Andric /// Bundle1
305bdd1243dSDimitry Andric /// ...
306bdd1243dSDimitry Andric /// BundleN
307bdd1243dSDimitry Andric 
308bdd1243dSDimitry Andric /// Read 8-byte integers from a buffer in little-endian format.
309bdd1243dSDimitry Andric static uint64_t Read8byteIntegerFromBuffer(StringRef Buffer, size_t pos) {
310bdd1243dSDimitry Andric   return llvm::support::endian::read64le(Buffer.data() + pos);
311bdd1243dSDimitry Andric }
312bdd1243dSDimitry Andric 
313bdd1243dSDimitry Andric /// Write 8-byte integers to a buffer in little-endian format.
314bdd1243dSDimitry Andric static void Write8byteIntegerToBuffer(raw_fd_ostream &OS, uint64_t Val) {
315bdd1243dSDimitry Andric   llvm::support::endian::write(OS, Val, llvm::support::little);
316bdd1243dSDimitry Andric }
317bdd1243dSDimitry Andric 
318bdd1243dSDimitry Andric class BinaryFileHandler final : public FileHandler {
319bdd1243dSDimitry Andric   /// Information about the bundles extracted from the header.
320bdd1243dSDimitry Andric   struct BinaryBundleInfo final : public BundleInfo {
321bdd1243dSDimitry Andric     /// Size of the bundle.
322bdd1243dSDimitry Andric     uint64_t Size = 0u;
323bdd1243dSDimitry Andric     /// Offset at which the bundle starts in the bundled file.
324bdd1243dSDimitry Andric     uint64_t Offset = 0u;
325bdd1243dSDimitry Andric 
326bdd1243dSDimitry Andric     BinaryBundleInfo() {}
327bdd1243dSDimitry Andric     BinaryBundleInfo(uint64_t Size, uint64_t Offset)
328bdd1243dSDimitry Andric         : Size(Size), Offset(Offset) {}
329bdd1243dSDimitry Andric   };
330bdd1243dSDimitry Andric 
331bdd1243dSDimitry Andric   /// Map between a triple and the corresponding bundle information.
332bdd1243dSDimitry Andric   StringMap<BinaryBundleInfo> BundlesInfo;
333bdd1243dSDimitry Andric 
334bdd1243dSDimitry Andric   /// Iterator for the bundle information that is being read.
335bdd1243dSDimitry Andric   StringMap<BinaryBundleInfo>::iterator CurBundleInfo;
336bdd1243dSDimitry Andric   StringMap<BinaryBundleInfo>::iterator NextBundleInfo;
337bdd1243dSDimitry Andric 
338bdd1243dSDimitry Andric   /// Current bundle target to be written.
339bdd1243dSDimitry Andric   std::string CurWriteBundleTarget;
340bdd1243dSDimitry Andric 
341bdd1243dSDimitry Andric   /// Configuration options and arrays for this bundler job
342bdd1243dSDimitry Andric   const OffloadBundlerConfig &BundlerConfig;
343bdd1243dSDimitry Andric 
344bdd1243dSDimitry Andric public:
345bdd1243dSDimitry Andric   // TODO: Add error checking from ClangOffloadBundler.cpp
346bdd1243dSDimitry Andric   BinaryFileHandler(const OffloadBundlerConfig &BC) : BundlerConfig(BC) {}
347bdd1243dSDimitry Andric 
348bdd1243dSDimitry Andric   ~BinaryFileHandler() final {}
349bdd1243dSDimitry Andric 
350bdd1243dSDimitry Andric   Error ReadHeader(MemoryBuffer &Input) final {
351bdd1243dSDimitry Andric     StringRef FC = Input.getBuffer();
352bdd1243dSDimitry Andric 
353bdd1243dSDimitry Andric     // Initialize the current bundle with the end of the container.
354bdd1243dSDimitry Andric     CurBundleInfo = BundlesInfo.end();
355bdd1243dSDimitry Andric 
356bdd1243dSDimitry Andric     // Check if buffer is smaller than magic string.
357bdd1243dSDimitry Andric     size_t ReadChars = sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1;
358bdd1243dSDimitry Andric     if (ReadChars > FC.size())
359bdd1243dSDimitry Andric       return Error::success();
360bdd1243dSDimitry Andric 
361bdd1243dSDimitry Andric     // Check if no magic was found.
362bdd1243dSDimitry Andric     StringRef Magic(FC.data(), sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1);
363bdd1243dSDimitry Andric     if (!Magic.equals(OFFLOAD_BUNDLER_MAGIC_STR))
364bdd1243dSDimitry Andric       return Error::success();
365bdd1243dSDimitry Andric 
366bdd1243dSDimitry Andric     // Read number of bundles.
367bdd1243dSDimitry Andric     if (ReadChars + 8 > FC.size())
368bdd1243dSDimitry Andric       return Error::success();
369bdd1243dSDimitry Andric 
370bdd1243dSDimitry Andric     uint64_t NumberOfBundles = Read8byteIntegerFromBuffer(FC, ReadChars);
371bdd1243dSDimitry Andric     ReadChars += 8;
372bdd1243dSDimitry Andric 
373bdd1243dSDimitry Andric     // Read bundle offsets, sizes and triples.
374bdd1243dSDimitry Andric     for (uint64_t i = 0; i < NumberOfBundles; ++i) {
375bdd1243dSDimitry Andric 
376bdd1243dSDimitry Andric       // Read offset.
377bdd1243dSDimitry Andric       if (ReadChars + 8 > FC.size())
378bdd1243dSDimitry Andric         return Error::success();
379bdd1243dSDimitry Andric 
380bdd1243dSDimitry Andric       uint64_t Offset = Read8byteIntegerFromBuffer(FC, ReadChars);
381bdd1243dSDimitry Andric       ReadChars += 8;
382bdd1243dSDimitry Andric 
383bdd1243dSDimitry Andric       // Read size.
384bdd1243dSDimitry Andric       if (ReadChars + 8 > FC.size())
385bdd1243dSDimitry Andric         return Error::success();
386bdd1243dSDimitry Andric 
387bdd1243dSDimitry Andric       uint64_t Size = Read8byteIntegerFromBuffer(FC, ReadChars);
388bdd1243dSDimitry Andric       ReadChars += 8;
389bdd1243dSDimitry Andric 
390bdd1243dSDimitry Andric       // Read triple size.
391bdd1243dSDimitry Andric       if (ReadChars + 8 > FC.size())
392bdd1243dSDimitry Andric         return Error::success();
393bdd1243dSDimitry Andric 
394bdd1243dSDimitry Andric       uint64_t TripleSize = Read8byteIntegerFromBuffer(FC, ReadChars);
395bdd1243dSDimitry Andric       ReadChars += 8;
396bdd1243dSDimitry Andric 
397bdd1243dSDimitry Andric       // Read triple.
398bdd1243dSDimitry Andric       if (ReadChars + TripleSize > FC.size())
399bdd1243dSDimitry Andric         return Error::success();
400bdd1243dSDimitry Andric 
401bdd1243dSDimitry Andric       StringRef Triple(&FC.data()[ReadChars], TripleSize);
402bdd1243dSDimitry Andric       ReadChars += TripleSize;
403bdd1243dSDimitry Andric 
404bdd1243dSDimitry Andric       // Check if the offset and size make sense.
405bdd1243dSDimitry Andric       if (!Offset || Offset + Size > FC.size())
406bdd1243dSDimitry Andric         return Error::success();
407bdd1243dSDimitry Andric 
408*06c3fb27SDimitry Andric       assert(!BundlesInfo.contains(Triple) && "Triple is duplicated??");
409bdd1243dSDimitry Andric       BundlesInfo[Triple] = BinaryBundleInfo(Size, Offset);
410bdd1243dSDimitry Andric     }
411bdd1243dSDimitry Andric     // Set the iterator to where we will start to read.
412bdd1243dSDimitry Andric     CurBundleInfo = BundlesInfo.end();
413bdd1243dSDimitry Andric     NextBundleInfo = BundlesInfo.begin();
414bdd1243dSDimitry Andric     return Error::success();
415bdd1243dSDimitry Andric   }
416bdd1243dSDimitry Andric 
417bdd1243dSDimitry Andric   Expected<std::optional<StringRef>>
418bdd1243dSDimitry Andric   ReadBundleStart(MemoryBuffer &Input) final {
419bdd1243dSDimitry Andric     if (NextBundleInfo == BundlesInfo.end())
420bdd1243dSDimitry Andric       return std::nullopt;
421bdd1243dSDimitry Andric     CurBundleInfo = NextBundleInfo++;
422bdd1243dSDimitry Andric     return CurBundleInfo->first();
423bdd1243dSDimitry Andric   }
424bdd1243dSDimitry Andric 
425bdd1243dSDimitry Andric   Error ReadBundleEnd(MemoryBuffer &Input) final {
426bdd1243dSDimitry Andric     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
427bdd1243dSDimitry Andric     return Error::success();
428bdd1243dSDimitry Andric   }
429bdd1243dSDimitry Andric 
430bdd1243dSDimitry Andric   Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
431bdd1243dSDimitry Andric     assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!");
432bdd1243dSDimitry Andric     StringRef FC = Input.getBuffer();
433bdd1243dSDimitry Andric     OS.write(FC.data() + CurBundleInfo->second.Offset,
434bdd1243dSDimitry Andric              CurBundleInfo->second.Size);
435bdd1243dSDimitry Andric     return Error::success();
436bdd1243dSDimitry Andric   }
437bdd1243dSDimitry Andric 
438bdd1243dSDimitry Andric   Error WriteHeader(raw_fd_ostream &OS,
439bdd1243dSDimitry Andric                     ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
440bdd1243dSDimitry Andric 
441bdd1243dSDimitry Andric     // Compute size of the header.
442bdd1243dSDimitry Andric     uint64_t HeaderSize = 0;
443bdd1243dSDimitry Andric 
444bdd1243dSDimitry Andric     HeaderSize += sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1;
445bdd1243dSDimitry Andric     HeaderSize += 8; // Number of Bundles
446bdd1243dSDimitry Andric 
447bdd1243dSDimitry Andric     for (auto &T : BundlerConfig.TargetNames) {
448bdd1243dSDimitry Andric       HeaderSize += 3 * 8; // Bundle offset, Size of bundle and size of triple.
449bdd1243dSDimitry Andric       HeaderSize += T.size(); // The triple.
450bdd1243dSDimitry Andric     }
451bdd1243dSDimitry Andric 
452bdd1243dSDimitry Andric     // Write to the buffer the header.
453bdd1243dSDimitry Andric     OS << OFFLOAD_BUNDLER_MAGIC_STR;
454bdd1243dSDimitry Andric 
455bdd1243dSDimitry Andric     Write8byteIntegerToBuffer(OS, BundlerConfig.TargetNames.size());
456bdd1243dSDimitry Andric 
457bdd1243dSDimitry Andric     unsigned Idx = 0;
458bdd1243dSDimitry Andric     for (auto &T : BundlerConfig.TargetNames) {
459bdd1243dSDimitry Andric       MemoryBuffer &MB = *Inputs[Idx++];
460bdd1243dSDimitry Andric       HeaderSize = alignTo(HeaderSize, BundlerConfig.BundleAlignment);
461bdd1243dSDimitry Andric       // Bundle offset.
462bdd1243dSDimitry Andric       Write8byteIntegerToBuffer(OS, HeaderSize);
463bdd1243dSDimitry Andric       // Size of the bundle (adds to the next bundle's offset)
464bdd1243dSDimitry Andric       Write8byteIntegerToBuffer(OS, MB.getBufferSize());
465bdd1243dSDimitry Andric       BundlesInfo[T] = BinaryBundleInfo(MB.getBufferSize(), HeaderSize);
466bdd1243dSDimitry Andric       HeaderSize += MB.getBufferSize();
467bdd1243dSDimitry Andric       // Size of the triple
468bdd1243dSDimitry Andric       Write8byteIntegerToBuffer(OS, T.size());
469bdd1243dSDimitry Andric       // Triple
470bdd1243dSDimitry Andric       OS << T;
471bdd1243dSDimitry Andric     }
472bdd1243dSDimitry Andric     return Error::success();
473bdd1243dSDimitry Andric   }
474bdd1243dSDimitry Andric 
475bdd1243dSDimitry Andric   Error WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) final {
476bdd1243dSDimitry Andric     CurWriteBundleTarget = TargetTriple.str();
477bdd1243dSDimitry Andric     return Error::success();
478bdd1243dSDimitry Andric   }
479bdd1243dSDimitry Andric 
480bdd1243dSDimitry Andric   Error WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) final {
481bdd1243dSDimitry Andric     return Error::success();
482bdd1243dSDimitry Andric   }
483bdd1243dSDimitry Andric 
484bdd1243dSDimitry Andric   Error WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
485bdd1243dSDimitry Andric     auto BI = BundlesInfo[CurWriteBundleTarget];
486bdd1243dSDimitry Andric     OS.seek(BI.Offset);
487bdd1243dSDimitry Andric     OS.write(Input.getBufferStart(), Input.getBufferSize());
488bdd1243dSDimitry Andric     return Error::success();
489bdd1243dSDimitry Andric   }
490bdd1243dSDimitry Andric };
491bdd1243dSDimitry Andric 
492bdd1243dSDimitry Andric // This class implements a list of temporary files that are removed upon
493bdd1243dSDimitry Andric // object destruction.
494bdd1243dSDimitry Andric class TempFileHandlerRAII {
495bdd1243dSDimitry Andric public:
496bdd1243dSDimitry Andric   ~TempFileHandlerRAII() {
497bdd1243dSDimitry Andric     for (const auto &File : Files)
498bdd1243dSDimitry Andric       sys::fs::remove(File);
499bdd1243dSDimitry Andric   }
500bdd1243dSDimitry Andric 
501bdd1243dSDimitry Andric   // Creates temporary file with given contents.
502bdd1243dSDimitry Andric   Expected<StringRef> Create(std::optional<ArrayRef<char>> Contents) {
503bdd1243dSDimitry Andric     SmallString<128u> File;
504bdd1243dSDimitry Andric     if (std::error_code EC =
505bdd1243dSDimitry Andric             sys::fs::createTemporaryFile("clang-offload-bundler", "tmp", File))
506bdd1243dSDimitry Andric       return createFileError(File, EC);
507bdd1243dSDimitry Andric     Files.push_front(File);
508bdd1243dSDimitry Andric 
509bdd1243dSDimitry Andric     if (Contents) {
510bdd1243dSDimitry Andric       std::error_code EC;
511bdd1243dSDimitry Andric       raw_fd_ostream OS(File, EC);
512bdd1243dSDimitry Andric       if (EC)
513bdd1243dSDimitry Andric         return createFileError(File, EC);
514bdd1243dSDimitry Andric       OS.write(Contents->data(), Contents->size());
515bdd1243dSDimitry Andric     }
516bdd1243dSDimitry Andric     return Files.front().str();
517bdd1243dSDimitry Andric   }
518bdd1243dSDimitry Andric 
519bdd1243dSDimitry Andric private:
520bdd1243dSDimitry Andric   std::forward_list<SmallString<128u>> Files;
521bdd1243dSDimitry Andric };
522bdd1243dSDimitry Andric 
523bdd1243dSDimitry Andric /// Handler for object files. The bundles are organized by sections with a
524bdd1243dSDimitry Andric /// designated name.
525bdd1243dSDimitry Andric ///
526bdd1243dSDimitry Andric /// To unbundle, we just copy the contents of the designated section.
527bdd1243dSDimitry Andric class ObjectFileHandler final : public FileHandler {
528bdd1243dSDimitry Andric 
529bdd1243dSDimitry Andric   /// The object file we are currently dealing with.
530bdd1243dSDimitry Andric   std::unique_ptr<ObjectFile> Obj;
531bdd1243dSDimitry Andric 
532bdd1243dSDimitry Andric   /// Return the input file contents.
533bdd1243dSDimitry Andric   StringRef getInputFileContents() const { return Obj->getData(); }
534bdd1243dSDimitry Andric 
535bdd1243dSDimitry Andric   /// Return bundle name (<kind>-<triple>) if the provided section is an offload
536bdd1243dSDimitry Andric   /// section.
537bdd1243dSDimitry Andric   static Expected<std::optional<StringRef>>
538bdd1243dSDimitry Andric   IsOffloadSection(SectionRef CurSection) {
539bdd1243dSDimitry Andric     Expected<StringRef> NameOrErr = CurSection.getName();
540bdd1243dSDimitry Andric     if (!NameOrErr)
541bdd1243dSDimitry Andric       return NameOrErr.takeError();
542bdd1243dSDimitry Andric 
543bdd1243dSDimitry Andric     // If it does not start with the reserved suffix, just skip this section.
544bdd1243dSDimitry Andric     if (!NameOrErr->startswith(OFFLOAD_BUNDLER_MAGIC_STR))
545bdd1243dSDimitry Andric       return std::nullopt;
546bdd1243dSDimitry Andric 
547bdd1243dSDimitry Andric     // Return the triple that is right after the reserved prefix.
548bdd1243dSDimitry Andric     return NameOrErr->substr(sizeof(OFFLOAD_BUNDLER_MAGIC_STR) - 1);
549bdd1243dSDimitry Andric   }
550bdd1243dSDimitry Andric 
551bdd1243dSDimitry Andric   /// Total number of inputs.
552bdd1243dSDimitry Andric   unsigned NumberOfInputs = 0;
553bdd1243dSDimitry Andric 
554bdd1243dSDimitry Andric   /// Total number of processed inputs, i.e, inputs that were already
555bdd1243dSDimitry Andric   /// read from the buffers.
556bdd1243dSDimitry Andric   unsigned NumberOfProcessedInputs = 0;
557bdd1243dSDimitry Andric 
558bdd1243dSDimitry Andric   /// Iterator of the current and next section.
559bdd1243dSDimitry Andric   section_iterator CurrentSection;
560bdd1243dSDimitry Andric   section_iterator NextSection;
561bdd1243dSDimitry Andric 
562bdd1243dSDimitry Andric   /// Configuration options and arrays for this bundler job
563bdd1243dSDimitry Andric   const OffloadBundlerConfig &BundlerConfig;
564bdd1243dSDimitry Andric 
565bdd1243dSDimitry Andric public:
566bdd1243dSDimitry Andric   // TODO: Add error checking from ClangOffloadBundler.cpp
567bdd1243dSDimitry Andric   ObjectFileHandler(std::unique_ptr<ObjectFile> ObjIn,
568bdd1243dSDimitry Andric                     const OffloadBundlerConfig &BC)
569bdd1243dSDimitry Andric       : Obj(std::move(ObjIn)), CurrentSection(Obj->section_begin()),
570bdd1243dSDimitry Andric         NextSection(Obj->section_begin()), BundlerConfig(BC) {}
571bdd1243dSDimitry Andric 
572bdd1243dSDimitry Andric   ~ObjectFileHandler() final {}
573bdd1243dSDimitry Andric 
574bdd1243dSDimitry Andric   Error ReadHeader(MemoryBuffer &Input) final { return Error::success(); }
575bdd1243dSDimitry Andric 
576bdd1243dSDimitry Andric   Expected<std::optional<StringRef>>
577bdd1243dSDimitry Andric   ReadBundleStart(MemoryBuffer &Input) final {
578bdd1243dSDimitry Andric     while (NextSection != Obj->section_end()) {
579bdd1243dSDimitry Andric       CurrentSection = NextSection;
580bdd1243dSDimitry Andric       ++NextSection;
581bdd1243dSDimitry Andric 
582bdd1243dSDimitry Andric       // Check if the current section name starts with the reserved prefix. If
583bdd1243dSDimitry Andric       // so, return the triple.
584bdd1243dSDimitry Andric       Expected<std::optional<StringRef>> TripleOrErr =
585bdd1243dSDimitry Andric           IsOffloadSection(*CurrentSection);
586bdd1243dSDimitry Andric       if (!TripleOrErr)
587bdd1243dSDimitry Andric         return TripleOrErr.takeError();
588bdd1243dSDimitry Andric       if (*TripleOrErr)
589bdd1243dSDimitry Andric         return **TripleOrErr;
590bdd1243dSDimitry Andric     }
591bdd1243dSDimitry Andric     return std::nullopt;
592bdd1243dSDimitry Andric   }
593bdd1243dSDimitry Andric 
594bdd1243dSDimitry Andric   Error ReadBundleEnd(MemoryBuffer &Input) final { return Error::success(); }
595bdd1243dSDimitry Andric 
596bdd1243dSDimitry Andric   Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
597bdd1243dSDimitry Andric     Expected<StringRef> ContentOrErr = CurrentSection->getContents();
598bdd1243dSDimitry Andric     if (!ContentOrErr)
599bdd1243dSDimitry Andric       return ContentOrErr.takeError();
600bdd1243dSDimitry Andric     StringRef Content = *ContentOrErr;
601bdd1243dSDimitry Andric 
602bdd1243dSDimitry Andric     // Copy fat object contents to the output when extracting host bundle.
603bdd1243dSDimitry Andric     if (Content.size() == 1u && Content.front() == 0)
604bdd1243dSDimitry Andric       Content = StringRef(Input.getBufferStart(), Input.getBufferSize());
605bdd1243dSDimitry Andric 
606bdd1243dSDimitry Andric     OS.write(Content.data(), Content.size());
607bdd1243dSDimitry Andric     return Error::success();
608bdd1243dSDimitry Andric   }
609bdd1243dSDimitry Andric 
610bdd1243dSDimitry Andric   Error WriteHeader(raw_fd_ostream &OS,
611bdd1243dSDimitry Andric                     ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
612bdd1243dSDimitry Andric     assert(BundlerConfig.HostInputIndex != ~0u &&
613bdd1243dSDimitry Andric            "Host input index not defined.");
614bdd1243dSDimitry Andric 
615bdd1243dSDimitry Andric     // Record number of inputs.
616bdd1243dSDimitry Andric     NumberOfInputs = Inputs.size();
617bdd1243dSDimitry Andric     return Error::success();
618bdd1243dSDimitry Andric   }
619bdd1243dSDimitry Andric 
620bdd1243dSDimitry Andric   Error WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) final {
621bdd1243dSDimitry Andric     ++NumberOfProcessedInputs;
622bdd1243dSDimitry Andric     return Error::success();
623bdd1243dSDimitry Andric   }
624bdd1243dSDimitry Andric 
625bdd1243dSDimitry Andric   Error WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) final {
626bdd1243dSDimitry Andric     assert(NumberOfProcessedInputs <= NumberOfInputs &&
627bdd1243dSDimitry Andric            "Processing more inputs that actually exist!");
628bdd1243dSDimitry Andric     assert(BundlerConfig.HostInputIndex != ~0u &&
629bdd1243dSDimitry Andric            "Host input index not defined.");
630bdd1243dSDimitry Andric 
631bdd1243dSDimitry Andric     // If this is not the last output, we don't have to do anything.
632bdd1243dSDimitry Andric     if (NumberOfProcessedInputs != NumberOfInputs)
633bdd1243dSDimitry Andric       return Error::success();
634bdd1243dSDimitry Andric 
635bdd1243dSDimitry Andric     // We will use llvm-objcopy to add target objects sections to the output
636bdd1243dSDimitry Andric     // fat object. These sections should have 'exclude' flag set which tells
637bdd1243dSDimitry Andric     // link editor to remove them from linker inputs when linking executable or
638bdd1243dSDimitry Andric     // shared library.
639bdd1243dSDimitry Andric 
640bdd1243dSDimitry Andric     assert(BundlerConfig.ObjcopyPath != "" &&
641bdd1243dSDimitry Andric            "llvm-objcopy path not specified");
642bdd1243dSDimitry Andric 
643bdd1243dSDimitry Andric     // We write to the output file directly. So, we close it and use the name
644bdd1243dSDimitry Andric     // to pass down to llvm-objcopy.
645bdd1243dSDimitry Andric     OS.close();
646bdd1243dSDimitry Andric 
647bdd1243dSDimitry Andric     // Temporary files that need to be removed.
648bdd1243dSDimitry Andric     TempFileHandlerRAII TempFiles;
649bdd1243dSDimitry Andric 
650bdd1243dSDimitry Andric     // Compose llvm-objcopy command line for add target objects' sections with
651bdd1243dSDimitry Andric     // appropriate flags.
652bdd1243dSDimitry Andric     BumpPtrAllocator Alloc;
653bdd1243dSDimitry Andric     StringSaver SS{Alloc};
654bdd1243dSDimitry Andric     SmallVector<StringRef, 8u> ObjcopyArgs{"llvm-objcopy"};
655bdd1243dSDimitry Andric 
656bdd1243dSDimitry Andric     for (unsigned I = 0; I < NumberOfInputs; ++I) {
657bdd1243dSDimitry Andric       StringRef InputFile = BundlerConfig.InputFileNames[I];
658bdd1243dSDimitry Andric       if (I == BundlerConfig.HostInputIndex) {
659bdd1243dSDimitry Andric         // Special handling for the host bundle. We do not need to add a
660bdd1243dSDimitry Andric         // standard bundle for the host object since we are going to use fat
661bdd1243dSDimitry Andric         // object as a host object. Therefore use dummy contents (one zero byte)
662bdd1243dSDimitry Andric         // when creating section for the host bundle.
663bdd1243dSDimitry Andric         Expected<StringRef> TempFileOrErr = TempFiles.Create(ArrayRef<char>(0));
664bdd1243dSDimitry Andric         if (!TempFileOrErr)
665bdd1243dSDimitry Andric           return TempFileOrErr.takeError();
666bdd1243dSDimitry Andric         InputFile = *TempFileOrErr;
667bdd1243dSDimitry Andric       }
668bdd1243dSDimitry Andric 
669bdd1243dSDimitry Andric       ObjcopyArgs.push_back(
670bdd1243dSDimitry Andric           SS.save(Twine("--add-section=") + OFFLOAD_BUNDLER_MAGIC_STR +
671bdd1243dSDimitry Andric                   BundlerConfig.TargetNames[I] + "=" + InputFile));
672bdd1243dSDimitry Andric       ObjcopyArgs.push_back(
673bdd1243dSDimitry Andric           SS.save(Twine("--set-section-flags=") + OFFLOAD_BUNDLER_MAGIC_STR +
674bdd1243dSDimitry Andric                   BundlerConfig.TargetNames[I] + "=readonly,exclude"));
675bdd1243dSDimitry Andric     }
676bdd1243dSDimitry Andric     ObjcopyArgs.push_back("--");
677bdd1243dSDimitry Andric     ObjcopyArgs.push_back(
678bdd1243dSDimitry Andric         BundlerConfig.InputFileNames[BundlerConfig.HostInputIndex]);
679bdd1243dSDimitry Andric     ObjcopyArgs.push_back(BundlerConfig.OutputFileNames.front());
680bdd1243dSDimitry Andric 
681bdd1243dSDimitry Andric     if (Error Err = executeObjcopy(BundlerConfig.ObjcopyPath, ObjcopyArgs))
682bdd1243dSDimitry Andric       return Err;
683bdd1243dSDimitry Andric 
684bdd1243dSDimitry Andric     return Error::success();
685bdd1243dSDimitry Andric   }
686bdd1243dSDimitry Andric 
687bdd1243dSDimitry Andric   Error WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
688bdd1243dSDimitry Andric     return Error::success();
689bdd1243dSDimitry Andric   }
690bdd1243dSDimitry Andric 
691bdd1243dSDimitry Andric private:
692bdd1243dSDimitry Andric   Error executeObjcopy(StringRef Objcopy, ArrayRef<StringRef> Args) {
693bdd1243dSDimitry Andric     // If the user asked for the commands to be printed out, we do that
694bdd1243dSDimitry Andric     // instead of executing it.
695bdd1243dSDimitry Andric     if (BundlerConfig.PrintExternalCommands) {
696bdd1243dSDimitry Andric       errs() << "\"" << Objcopy << "\"";
697bdd1243dSDimitry Andric       for (StringRef Arg : drop_begin(Args, 1))
698bdd1243dSDimitry Andric         errs() << " \"" << Arg << "\"";
699bdd1243dSDimitry Andric       errs() << "\n";
700bdd1243dSDimitry Andric     } else {
701bdd1243dSDimitry Andric       if (sys::ExecuteAndWait(Objcopy, Args))
702bdd1243dSDimitry Andric         return createStringError(inconvertibleErrorCode(),
703bdd1243dSDimitry Andric                                  "'llvm-objcopy' tool failed");
704bdd1243dSDimitry Andric     }
705bdd1243dSDimitry Andric     return Error::success();
706bdd1243dSDimitry Andric   }
707bdd1243dSDimitry Andric };
708bdd1243dSDimitry Andric 
709bdd1243dSDimitry Andric /// Handler for text files. The bundled file will have the following format.
710bdd1243dSDimitry Andric ///
711bdd1243dSDimitry Andric /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
712bdd1243dSDimitry Andric /// Bundle 1
713bdd1243dSDimitry Andric /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__END__ triple"
714bdd1243dSDimitry Andric /// ...
715bdd1243dSDimitry Andric /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__START__ triple"
716bdd1243dSDimitry Andric /// Bundle N
717bdd1243dSDimitry Andric /// "Comment OFFLOAD_BUNDLER_MAGIC_STR__END__ triple"
718bdd1243dSDimitry Andric class TextFileHandler final : public FileHandler {
719bdd1243dSDimitry Andric   /// String that begins a line comment.
720bdd1243dSDimitry Andric   StringRef Comment;
721bdd1243dSDimitry Andric 
722bdd1243dSDimitry Andric   /// String that initiates a bundle.
723bdd1243dSDimitry Andric   std::string BundleStartString;
724bdd1243dSDimitry Andric 
725bdd1243dSDimitry Andric   /// String that closes a bundle.
726bdd1243dSDimitry Andric   std::string BundleEndString;
727bdd1243dSDimitry Andric 
728bdd1243dSDimitry Andric   /// Number of chars read from input.
729bdd1243dSDimitry Andric   size_t ReadChars = 0u;
730bdd1243dSDimitry Andric 
731bdd1243dSDimitry Andric protected:
732bdd1243dSDimitry Andric   Error ReadHeader(MemoryBuffer &Input) final { return Error::success(); }
733bdd1243dSDimitry Andric 
734bdd1243dSDimitry Andric   Expected<std::optional<StringRef>>
735bdd1243dSDimitry Andric   ReadBundleStart(MemoryBuffer &Input) final {
736bdd1243dSDimitry Andric     StringRef FC = Input.getBuffer();
737bdd1243dSDimitry Andric 
738bdd1243dSDimitry Andric     // Find start of the bundle.
739bdd1243dSDimitry Andric     ReadChars = FC.find(BundleStartString, ReadChars);
740bdd1243dSDimitry Andric     if (ReadChars == FC.npos)
741bdd1243dSDimitry Andric       return std::nullopt;
742bdd1243dSDimitry Andric 
743bdd1243dSDimitry Andric     // Get position of the triple.
744bdd1243dSDimitry Andric     size_t TripleStart = ReadChars = ReadChars + BundleStartString.size();
745bdd1243dSDimitry Andric 
746bdd1243dSDimitry Andric     // Get position that closes the triple.
747bdd1243dSDimitry Andric     size_t TripleEnd = ReadChars = FC.find("\n", ReadChars);
748bdd1243dSDimitry Andric     if (TripleEnd == FC.npos)
749bdd1243dSDimitry Andric       return std::nullopt;
750bdd1243dSDimitry Andric 
751bdd1243dSDimitry Andric     // Next time we read after the new line.
752bdd1243dSDimitry Andric     ++ReadChars;
753bdd1243dSDimitry Andric 
754bdd1243dSDimitry Andric     return StringRef(&FC.data()[TripleStart], TripleEnd - TripleStart);
755bdd1243dSDimitry Andric   }
756bdd1243dSDimitry Andric 
757bdd1243dSDimitry Andric   Error ReadBundleEnd(MemoryBuffer &Input) final {
758bdd1243dSDimitry Andric     StringRef FC = Input.getBuffer();
759bdd1243dSDimitry Andric 
760bdd1243dSDimitry Andric     // Read up to the next new line.
761bdd1243dSDimitry Andric     assert(FC[ReadChars] == '\n' && "The bundle should end with a new line.");
762bdd1243dSDimitry Andric 
763bdd1243dSDimitry Andric     size_t TripleEnd = ReadChars = FC.find("\n", ReadChars + 1);
764bdd1243dSDimitry Andric     if (TripleEnd != FC.npos)
765bdd1243dSDimitry Andric       // Next time we read after the new line.
766bdd1243dSDimitry Andric       ++ReadChars;
767bdd1243dSDimitry Andric 
768bdd1243dSDimitry Andric     return Error::success();
769bdd1243dSDimitry Andric   }
770bdd1243dSDimitry Andric 
771bdd1243dSDimitry Andric   Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final {
772bdd1243dSDimitry Andric     StringRef FC = Input.getBuffer();
773bdd1243dSDimitry Andric     size_t BundleStart = ReadChars;
774bdd1243dSDimitry Andric 
775bdd1243dSDimitry Andric     // Find end of the bundle.
776bdd1243dSDimitry Andric     size_t BundleEnd = ReadChars = FC.find(BundleEndString, ReadChars);
777bdd1243dSDimitry Andric 
778bdd1243dSDimitry Andric     StringRef Bundle(&FC.data()[BundleStart], BundleEnd - BundleStart);
779bdd1243dSDimitry Andric     OS << Bundle;
780bdd1243dSDimitry Andric 
781bdd1243dSDimitry Andric     return Error::success();
782bdd1243dSDimitry Andric   }
783bdd1243dSDimitry Andric 
784bdd1243dSDimitry Andric   Error WriteHeader(raw_fd_ostream &OS,
785bdd1243dSDimitry Andric                     ArrayRef<std::unique_ptr<MemoryBuffer>> Inputs) final {
786bdd1243dSDimitry Andric     return Error::success();
787bdd1243dSDimitry Andric   }
788bdd1243dSDimitry Andric 
789bdd1243dSDimitry Andric   Error WriteBundleStart(raw_fd_ostream &OS, StringRef TargetTriple) final {
790bdd1243dSDimitry Andric     OS << BundleStartString << TargetTriple << "\n";
791bdd1243dSDimitry Andric     return Error::success();
792bdd1243dSDimitry Andric   }
793bdd1243dSDimitry Andric 
794bdd1243dSDimitry Andric   Error WriteBundleEnd(raw_fd_ostream &OS, StringRef TargetTriple) final {
795bdd1243dSDimitry Andric     OS << BundleEndString << TargetTriple << "\n";
796bdd1243dSDimitry Andric     return Error::success();
797bdd1243dSDimitry Andric   }
798bdd1243dSDimitry Andric 
799bdd1243dSDimitry Andric   Error WriteBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final {
800bdd1243dSDimitry Andric     OS << Input.getBuffer();
801bdd1243dSDimitry Andric     return Error::success();
802bdd1243dSDimitry Andric   }
803bdd1243dSDimitry Andric 
804bdd1243dSDimitry Andric public:
805bdd1243dSDimitry Andric   TextFileHandler(StringRef Comment) : Comment(Comment), ReadChars(0) {
806bdd1243dSDimitry Andric     BundleStartString =
807bdd1243dSDimitry Andric         "\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__START__ ";
808bdd1243dSDimitry Andric     BundleEndString =
809bdd1243dSDimitry Andric         "\n" + Comment.str() + " " OFFLOAD_BUNDLER_MAGIC_STR "__END__ ";
810bdd1243dSDimitry Andric   }
811bdd1243dSDimitry Andric 
812bdd1243dSDimitry Andric   Error listBundleIDsCallback(MemoryBuffer &Input,
813bdd1243dSDimitry Andric                               const BundleInfo &Info) final {
814bdd1243dSDimitry Andric     // TODO: To list bundle IDs in a bundled text file we need to go through
815bdd1243dSDimitry Andric     // all bundles. The format of bundled text file may need to include a
816bdd1243dSDimitry Andric     // header if the performance of listing bundle IDs of bundled text file is
817bdd1243dSDimitry Andric     // important.
818bdd1243dSDimitry Andric     ReadChars = Input.getBuffer().find(BundleEndString, ReadChars);
819bdd1243dSDimitry Andric     if (Error Err = ReadBundleEnd(Input))
820bdd1243dSDimitry Andric       return Err;
821bdd1243dSDimitry Andric     return Error::success();
822bdd1243dSDimitry Andric   }
823bdd1243dSDimitry Andric };
824bdd1243dSDimitry Andric } // namespace
825bdd1243dSDimitry Andric 
826bdd1243dSDimitry Andric /// Return an appropriate object file handler. We use the specific object
827bdd1243dSDimitry Andric /// handler if we know how to deal with that format, otherwise we use a default
828bdd1243dSDimitry Andric /// binary file handler.
829bdd1243dSDimitry Andric static std::unique_ptr<FileHandler>
830bdd1243dSDimitry Andric CreateObjectFileHandler(MemoryBuffer &FirstInput,
831bdd1243dSDimitry Andric                         const OffloadBundlerConfig &BundlerConfig) {
832bdd1243dSDimitry Andric   // Check if the input file format is one that we know how to deal with.
833bdd1243dSDimitry Andric   Expected<std::unique_ptr<Binary>> BinaryOrErr = createBinary(FirstInput);
834bdd1243dSDimitry Andric 
835bdd1243dSDimitry Andric   // We only support regular object files. If failed to open the input as a
836bdd1243dSDimitry Andric   // known binary or this is not an object file use the default binary handler.
837bdd1243dSDimitry Andric   if (errorToBool(BinaryOrErr.takeError()) || !isa<ObjectFile>(*BinaryOrErr))
838bdd1243dSDimitry Andric     return std::make_unique<BinaryFileHandler>(BundlerConfig);
839bdd1243dSDimitry Andric 
840bdd1243dSDimitry Andric   // Otherwise create an object file handler. The handler will be owned by the
841bdd1243dSDimitry Andric   // client of this function.
842bdd1243dSDimitry Andric   return std::make_unique<ObjectFileHandler>(
843bdd1243dSDimitry Andric       std::unique_ptr<ObjectFile>(cast<ObjectFile>(BinaryOrErr->release())),
844bdd1243dSDimitry Andric       BundlerConfig);
845bdd1243dSDimitry Andric }
846bdd1243dSDimitry Andric 
847bdd1243dSDimitry Andric /// Return an appropriate handler given the input files and options.
848bdd1243dSDimitry Andric static Expected<std::unique_ptr<FileHandler>>
849bdd1243dSDimitry Andric CreateFileHandler(MemoryBuffer &FirstInput,
850bdd1243dSDimitry Andric                   const OffloadBundlerConfig &BundlerConfig) {
851bdd1243dSDimitry Andric   std::string FilesType = BundlerConfig.FilesType;
852bdd1243dSDimitry Andric 
853bdd1243dSDimitry Andric   if (FilesType == "i")
854bdd1243dSDimitry Andric     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
855bdd1243dSDimitry Andric   if (FilesType == "ii")
856bdd1243dSDimitry Andric     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
857bdd1243dSDimitry Andric   if (FilesType == "cui")
858bdd1243dSDimitry Andric     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
859bdd1243dSDimitry Andric   if (FilesType == "hipi")
860bdd1243dSDimitry Andric     return std::make_unique<TextFileHandler>(/*Comment=*/"//");
861bdd1243dSDimitry Andric   // TODO: `.d` should be eventually removed once `-M` and its variants are
862bdd1243dSDimitry Andric   // handled properly in offload compilation.
863bdd1243dSDimitry Andric   if (FilesType == "d")
864bdd1243dSDimitry Andric     return std::make_unique<TextFileHandler>(/*Comment=*/"#");
865bdd1243dSDimitry Andric   if (FilesType == "ll")
866bdd1243dSDimitry Andric     return std::make_unique<TextFileHandler>(/*Comment=*/";");
867bdd1243dSDimitry Andric   if (FilesType == "bc")
868bdd1243dSDimitry Andric     return std::make_unique<BinaryFileHandler>(BundlerConfig);
869bdd1243dSDimitry Andric   if (FilesType == "s")
870bdd1243dSDimitry Andric     return std::make_unique<TextFileHandler>(/*Comment=*/"#");
871bdd1243dSDimitry Andric   if (FilesType == "o")
872bdd1243dSDimitry Andric     return CreateObjectFileHandler(FirstInput, BundlerConfig);
873bdd1243dSDimitry Andric   if (FilesType == "a")
874bdd1243dSDimitry Andric     return CreateObjectFileHandler(FirstInput, BundlerConfig);
875bdd1243dSDimitry Andric   if (FilesType == "gch")
876bdd1243dSDimitry Andric     return std::make_unique<BinaryFileHandler>(BundlerConfig);
877bdd1243dSDimitry Andric   if (FilesType == "ast")
878bdd1243dSDimitry Andric     return std::make_unique<BinaryFileHandler>(BundlerConfig);
879bdd1243dSDimitry Andric 
880bdd1243dSDimitry Andric   return createStringError(errc::invalid_argument,
881bdd1243dSDimitry Andric                            "'" + FilesType + "': invalid file type specified");
882bdd1243dSDimitry Andric }
883bdd1243dSDimitry Andric 
884bdd1243dSDimitry Andric // List bundle IDs. Return true if an error was found.
885bdd1243dSDimitry Andric Error OffloadBundler::ListBundleIDsInFile(
886bdd1243dSDimitry Andric     StringRef InputFileName, const OffloadBundlerConfig &BundlerConfig) {
887bdd1243dSDimitry Andric   // Open Input file.
888bdd1243dSDimitry Andric   ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
889bdd1243dSDimitry Andric       MemoryBuffer::getFileOrSTDIN(InputFileName);
890bdd1243dSDimitry Andric   if (std::error_code EC = CodeOrErr.getError())
891bdd1243dSDimitry Andric     return createFileError(InputFileName, EC);
892bdd1243dSDimitry Andric 
893bdd1243dSDimitry Andric   MemoryBuffer &Input = **CodeOrErr;
894bdd1243dSDimitry Andric 
895bdd1243dSDimitry Andric   // Select the right files handler.
896bdd1243dSDimitry Andric   Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
897bdd1243dSDimitry Andric       CreateFileHandler(Input, BundlerConfig);
898bdd1243dSDimitry Andric   if (!FileHandlerOrErr)
899bdd1243dSDimitry Andric     return FileHandlerOrErr.takeError();
900bdd1243dSDimitry Andric 
901bdd1243dSDimitry Andric   std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
902bdd1243dSDimitry Andric   assert(FH);
903bdd1243dSDimitry Andric   return FH->listBundleIDs(Input);
904bdd1243dSDimitry Andric }
905bdd1243dSDimitry Andric 
906bdd1243dSDimitry Andric /// Bundle the files. Return true if an error was found.
907bdd1243dSDimitry Andric Error OffloadBundler::BundleFiles() {
908bdd1243dSDimitry Andric   std::error_code EC;
909bdd1243dSDimitry Andric 
910bdd1243dSDimitry Andric   // Create output file.
911bdd1243dSDimitry Andric   raw_fd_ostream OutputFile(BundlerConfig.OutputFileNames.front(), EC,
912bdd1243dSDimitry Andric                             sys::fs::OF_None);
913bdd1243dSDimitry Andric   if (EC)
914bdd1243dSDimitry Andric     return createFileError(BundlerConfig.OutputFileNames.front(), EC);
915bdd1243dSDimitry Andric 
916bdd1243dSDimitry Andric   // Open input files.
917bdd1243dSDimitry Andric   SmallVector<std::unique_ptr<MemoryBuffer>, 8u> InputBuffers;
918bdd1243dSDimitry Andric   InputBuffers.reserve(BundlerConfig.InputFileNames.size());
919bdd1243dSDimitry Andric   for (auto &I : BundlerConfig.InputFileNames) {
920bdd1243dSDimitry Andric     ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
921bdd1243dSDimitry Andric         MemoryBuffer::getFileOrSTDIN(I);
922bdd1243dSDimitry Andric     if (std::error_code EC = CodeOrErr.getError())
923bdd1243dSDimitry Andric       return createFileError(I, EC);
924bdd1243dSDimitry Andric     InputBuffers.emplace_back(std::move(*CodeOrErr));
925bdd1243dSDimitry Andric   }
926bdd1243dSDimitry Andric 
927bdd1243dSDimitry Andric   // Get the file handler. We use the host buffer as reference.
928bdd1243dSDimitry Andric   assert((BundlerConfig.HostInputIndex != ~0u || BundlerConfig.AllowNoHost) &&
929bdd1243dSDimitry Andric          "Host input index undefined??");
930bdd1243dSDimitry Andric   Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr = CreateFileHandler(
931bdd1243dSDimitry Andric       *InputBuffers[BundlerConfig.AllowNoHost ? 0
932bdd1243dSDimitry Andric                                               : BundlerConfig.HostInputIndex],
933bdd1243dSDimitry Andric       BundlerConfig);
934bdd1243dSDimitry Andric   if (!FileHandlerOrErr)
935bdd1243dSDimitry Andric     return FileHandlerOrErr.takeError();
936bdd1243dSDimitry Andric 
937bdd1243dSDimitry Andric   std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
938bdd1243dSDimitry Andric   assert(FH);
939bdd1243dSDimitry Andric 
940bdd1243dSDimitry Andric   // Write header.
941bdd1243dSDimitry Andric   if (Error Err = FH->WriteHeader(OutputFile, InputBuffers))
942bdd1243dSDimitry Andric     return Err;
943bdd1243dSDimitry Andric 
944bdd1243dSDimitry Andric   // Write all bundles along with the start/end markers. If an error was found
945bdd1243dSDimitry Andric   // writing the end of the bundle component, abort the bundle writing.
946bdd1243dSDimitry Andric   auto Input = InputBuffers.begin();
947bdd1243dSDimitry Andric   for (auto &Triple : BundlerConfig.TargetNames) {
948bdd1243dSDimitry Andric     if (Error Err = FH->WriteBundleStart(OutputFile, Triple))
949bdd1243dSDimitry Andric       return Err;
950bdd1243dSDimitry Andric     if (Error Err = FH->WriteBundle(OutputFile, **Input))
951bdd1243dSDimitry Andric       return Err;
952bdd1243dSDimitry Andric     if (Error Err = FH->WriteBundleEnd(OutputFile, Triple))
953bdd1243dSDimitry Andric       return Err;
954bdd1243dSDimitry Andric     ++Input;
955bdd1243dSDimitry Andric   }
956bdd1243dSDimitry Andric   return Error::success();
957bdd1243dSDimitry Andric }
958bdd1243dSDimitry Andric 
959bdd1243dSDimitry Andric // Unbundle the files. Return true if an error was found.
960bdd1243dSDimitry Andric Error OffloadBundler::UnbundleFiles() {
961bdd1243dSDimitry Andric   // Open Input file.
962bdd1243dSDimitry Andric   ErrorOr<std::unique_ptr<MemoryBuffer>> CodeOrErr =
963bdd1243dSDimitry Andric       MemoryBuffer::getFileOrSTDIN(BundlerConfig.InputFileNames.front());
964bdd1243dSDimitry Andric   if (std::error_code EC = CodeOrErr.getError())
965bdd1243dSDimitry Andric     return createFileError(BundlerConfig.InputFileNames.front(), EC);
966bdd1243dSDimitry Andric 
967bdd1243dSDimitry Andric   MemoryBuffer &Input = **CodeOrErr;
968bdd1243dSDimitry Andric 
969bdd1243dSDimitry Andric   // Select the right files handler.
970bdd1243dSDimitry Andric   Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
971bdd1243dSDimitry Andric       CreateFileHandler(Input, BundlerConfig);
972bdd1243dSDimitry Andric   if (!FileHandlerOrErr)
973bdd1243dSDimitry Andric     return FileHandlerOrErr.takeError();
974bdd1243dSDimitry Andric 
975bdd1243dSDimitry Andric   std::unique_ptr<FileHandler> &FH = *FileHandlerOrErr;
976bdd1243dSDimitry Andric   assert(FH);
977bdd1243dSDimitry Andric 
978bdd1243dSDimitry Andric   // Read the header of the bundled file.
979bdd1243dSDimitry Andric   if (Error Err = FH->ReadHeader(Input))
980bdd1243dSDimitry Andric     return Err;
981bdd1243dSDimitry Andric 
982bdd1243dSDimitry Andric   // Create a work list that consist of the map triple/output file.
983bdd1243dSDimitry Andric   StringMap<StringRef> Worklist;
984bdd1243dSDimitry Andric   auto Output = BundlerConfig.OutputFileNames.begin();
985bdd1243dSDimitry Andric   for (auto &Triple : BundlerConfig.TargetNames) {
986bdd1243dSDimitry Andric     Worklist[Triple] = *Output;
987bdd1243dSDimitry Andric     ++Output;
988bdd1243dSDimitry Andric   }
989bdd1243dSDimitry Andric 
990bdd1243dSDimitry Andric   // Read all the bundles that are in the work list. If we find no bundles we
991bdd1243dSDimitry Andric   // assume the file is meant for the host target.
992bdd1243dSDimitry Andric   bool FoundHostBundle = false;
993bdd1243dSDimitry Andric   while (!Worklist.empty()) {
994bdd1243dSDimitry Andric     Expected<std::optional<StringRef>> CurTripleOrErr =
995bdd1243dSDimitry Andric         FH->ReadBundleStart(Input);
996bdd1243dSDimitry Andric     if (!CurTripleOrErr)
997bdd1243dSDimitry Andric       return CurTripleOrErr.takeError();
998bdd1243dSDimitry Andric 
999bdd1243dSDimitry Andric     // We don't have more bundles.
1000bdd1243dSDimitry Andric     if (!*CurTripleOrErr)
1001bdd1243dSDimitry Andric       break;
1002bdd1243dSDimitry Andric 
1003bdd1243dSDimitry Andric     StringRef CurTriple = **CurTripleOrErr;
1004bdd1243dSDimitry Andric     assert(!CurTriple.empty());
1005bdd1243dSDimitry Andric 
1006bdd1243dSDimitry Andric     auto Output = Worklist.begin();
1007bdd1243dSDimitry Andric     for (auto E = Worklist.end(); Output != E; Output++) {
1008bdd1243dSDimitry Andric       if (isCodeObjectCompatible(
1009bdd1243dSDimitry Andric               OffloadTargetInfo(CurTriple, BundlerConfig),
1010bdd1243dSDimitry Andric               OffloadTargetInfo((*Output).first(), BundlerConfig))) {
1011bdd1243dSDimitry Andric         break;
1012bdd1243dSDimitry Andric       }
1013bdd1243dSDimitry Andric     }
1014bdd1243dSDimitry Andric 
1015bdd1243dSDimitry Andric     if (Output == Worklist.end())
1016bdd1243dSDimitry Andric       continue;
1017bdd1243dSDimitry Andric     // Check if the output file can be opened and copy the bundle to it.
1018bdd1243dSDimitry Andric     std::error_code EC;
1019bdd1243dSDimitry Andric     raw_fd_ostream OutputFile((*Output).second, EC, sys::fs::OF_None);
1020bdd1243dSDimitry Andric     if (EC)
1021bdd1243dSDimitry Andric       return createFileError((*Output).second, EC);
1022bdd1243dSDimitry Andric     if (Error Err = FH->ReadBundle(OutputFile, Input))
1023bdd1243dSDimitry Andric       return Err;
1024bdd1243dSDimitry Andric     if (Error Err = FH->ReadBundleEnd(Input))
1025bdd1243dSDimitry Andric       return Err;
1026bdd1243dSDimitry Andric     Worklist.erase(Output);
1027bdd1243dSDimitry Andric 
1028bdd1243dSDimitry Andric     // Record if we found the host bundle.
1029bdd1243dSDimitry Andric     auto OffloadInfo = OffloadTargetInfo(CurTriple, BundlerConfig);
1030bdd1243dSDimitry Andric     if (OffloadInfo.hasHostKind())
1031bdd1243dSDimitry Andric       FoundHostBundle = true;
1032bdd1243dSDimitry Andric   }
1033bdd1243dSDimitry Andric 
1034bdd1243dSDimitry Andric   if (!BundlerConfig.AllowMissingBundles && !Worklist.empty()) {
1035bdd1243dSDimitry Andric     std::string ErrMsg = "Can't find bundles for";
1036bdd1243dSDimitry Andric     std::set<StringRef> Sorted;
1037bdd1243dSDimitry Andric     for (auto &E : Worklist)
1038bdd1243dSDimitry Andric       Sorted.insert(E.first());
1039bdd1243dSDimitry Andric     unsigned I = 0;
1040bdd1243dSDimitry Andric     unsigned Last = Sorted.size() - 1;
1041bdd1243dSDimitry Andric     for (auto &E : Sorted) {
1042bdd1243dSDimitry Andric       if (I != 0 && Last > 1)
1043bdd1243dSDimitry Andric         ErrMsg += ",";
1044bdd1243dSDimitry Andric       ErrMsg += " ";
1045bdd1243dSDimitry Andric       if (I == Last && I != 0)
1046bdd1243dSDimitry Andric         ErrMsg += "and ";
1047bdd1243dSDimitry Andric       ErrMsg += E.str();
1048bdd1243dSDimitry Andric       ++I;
1049bdd1243dSDimitry Andric     }
1050bdd1243dSDimitry Andric     return createStringError(inconvertibleErrorCode(), ErrMsg);
1051bdd1243dSDimitry Andric   }
1052bdd1243dSDimitry Andric 
1053bdd1243dSDimitry Andric   // If no bundles were found, assume the input file is the host bundle and
1054bdd1243dSDimitry Andric   // create empty files for the remaining targets.
1055bdd1243dSDimitry Andric   if (Worklist.size() == BundlerConfig.TargetNames.size()) {
1056bdd1243dSDimitry Andric     for (auto &E : Worklist) {
1057bdd1243dSDimitry Andric       std::error_code EC;
1058bdd1243dSDimitry Andric       raw_fd_ostream OutputFile(E.second, EC, sys::fs::OF_None);
1059bdd1243dSDimitry Andric       if (EC)
1060bdd1243dSDimitry Andric         return createFileError(E.second, EC);
1061bdd1243dSDimitry Andric 
1062bdd1243dSDimitry Andric       // If this entry has a host kind, copy the input file to the output file.
1063bdd1243dSDimitry Andric       auto OffloadInfo = OffloadTargetInfo(E.getKey(), BundlerConfig);
1064bdd1243dSDimitry Andric       if (OffloadInfo.hasHostKind())
1065bdd1243dSDimitry Andric         OutputFile.write(Input.getBufferStart(), Input.getBufferSize());
1066bdd1243dSDimitry Andric     }
1067bdd1243dSDimitry Andric     return Error::success();
1068bdd1243dSDimitry Andric   }
1069bdd1243dSDimitry Andric 
1070bdd1243dSDimitry Andric   // If we found elements, we emit an error if none of those were for the host
1071bdd1243dSDimitry Andric   // in case host bundle name was provided in command line.
1072bdd1243dSDimitry Andric   if (!(FoundHostBundle || BundlerConfig.HostInputIndex == ~0u ||
1073bdd1243dSDimitry Andric         BundlerConfig.AllowMissingBundles))
1074bdd1243dSDimitry Andric     return createStringError(inconvertibleErrorCode(),
1075bdd1243dSDimitry Andric                              "Can't find bundle for the host target");
1076bdd1243dSDimitry Andric 
1077bdd1243dSDimitry Andric   // If we still have any elements in the worklist, create empty files for them.
1078bdd1243dSDimitry Andric   for (auto &E : Worklist) {
1079bdd1243dSDimitry Andric     std::error_code EC;
1080bdd1243dSDimitry Andric     raw_fd_ostream OutputFile(E.second, EC, sys::fs::OF_None);
1081bdd1243dSDimitry Andric     if (EC)
1082bdd1243dSDimitry Andric       return createFileError(E.second, EC);
1083bdd1243dSDimitry Andric   }
1084bdd1243dSDimitry Andric 
1085bdd1243dSDimitry Andric   return Error::success();
1086bdd1243dSDimitry Andric }
1087bdd1243dSDimitry Andric 
1088bdd1243dSDimitry Andric static Archive::Kind getDefaultArchiveKindForHost() {
1089bdd1243dSDimitry Andric   return Triple(sys::getDefaultTargetTriple()).isOSDarwin() ? Archive::K_DARWIN
1090bdd1243dSDimitry Andric                                                             : Archive::K_GNU;
1091bdd1243dSDimitry Andric }
1092bdd1243dSDimitry Andric 
1093bdd1243dSDimitry Andric /// @brief Computes a list of targets among all given targets which are
1094bdd1243dSDimitry Andric /// compatible with this code object
1095bdd1243dSDimitry Andric /// @param [in] CodeObjectInfo Code Object
1096bdd1243dSDimitry Andric /// @param [out] CompatibleTargets List of all compatible targets among all
1097bdd1243dSDimitry Andric /// given targets
1098bdd1243dSDimitry Andric /// @return false, if no compatible target is found.
1099bdd1243dSDimitry Andric static bool
1100bdd1243dSDimitry Andric getCompatibleOffloadTargets(OffloadTargetInfo &CodeObjectInfo,
1101bdd1243dSDimitry Andric                             SmallVectorImpl<StringRef> &CompatibleTargets,
1102bdd1243dSDimitry Andric                             const OffloadBundlerConfig &BundlerConfig) {
1103bdd1243dSDimitry Andric   if (!CompatibleTargets.empty()) {
1104bdd1243dSDimitry Andric     DEBUG_WITH_TYPE("CodeObjectCompatibility",
1105bdd1243dSDimitry Andric                     dbgs() << "CompatibleTargets list should be empty\n");
1106bdd1243dSDimitry Andric     return false;
1107bdd1243dSDimitry Andric   }
1108bdd1243dSDimitry Andric   for (auto &Target : BundlerConfig.TargetNames) {
1109bdd1243dSDimitry Andric     auto TargetInfo = OffloadTargetInfo(Target, BundlerConfig);
1110bdd1243dSDimitry Andric     if (isCodeObjectCompatible(CodeObjectInfo, TargetInfo))
1111bdd1243dSDimitry Andric       CompatibleTargets.push_back(Target);
1112bdd1243dSDimitry Andric   }
1113bdd1243dSDimitry Andric   return !CompatibleTargets.empty();
1114bdd1243dSDimitry Andric }
1115bdd1243dSDimitry Andric 
1116bdd1243dSDimitry Andric /// UnbundleArchive takes an archive file (".a") as input containing bundled
1117bdd1243dSDimitry Andric /// code object files, and a list of offload targets (not host), and extracts
1118bdd1243dSDimitry Andric /// the code objects into a new archive file for each offload target. Each
1119bdd1243dSDimitry Andric /// resulting archive file contains all code object files corresponding to that
1120bdd1243dSDimitry Andric /// particular offload target. The created archive file does not
1121bdd1243dSDimitry Andric /// contain an index of the symbols and code object files are named as
1122bdd1243dSDimitry Andric /// <<Parent Bundle Name>-<CodeObject's GPUArch>>, with ':' replaced with '_'.
1123bdd1243dSDimitry Andric Error OffloadBundler::UnbundleArchive() {
1124bdd1243dSDimitry Andric   std::vector<std::unique_ptr<MemoryBuffer>> ArchiveBuffers;
1125bdd1243dSDimitry Andric 
1126bdd1243dSDimitry Andric   /// Map of target names with list of object files that will form the device
1127bdd1243dSDimitry Andric   /// specific archive for that target
1128bdd1243dSDimitry Andric   StringMap<std::vector<NewArchiveMember>> OutputArchivesMap;
1129bdd1243dSDimitry Andric 
1130bdd1243dSDimitry Andric   // Map of target names and output archive filenames
1131bdd1243dSDimitry Andric   StringMap<StringRef> TargetOutputFileNameMap;
1132bdd1243dSDimitry Andric 
1133bdd1243dSDimitry Andric   auto Output = BundlerConfig.OutputFileNames.begin();
1134bdd1243dSDimitry Andric   for (auto &Target : BundlerConfig.TargetNames) {
1135bdd1243dSDimitry Andric     TargetOutputFileNameMap[Target] = *Output;
1136bdd1243dSDimitry Andric     ++Output;
1137bdd1243dSDimitry Andric   }
1138bdd1243dSDimitry Andric 
1139bdd1243dSDimitry Andric   StringRef IFName = BundlerConfig.InputFileNames.front();
1140bdd1243dSDimitry Andric 
1141bdd1243dSDimitry Andric   ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
1142bdd1243dSDimitry Andric       MemoryBuffer::getFileOrSTDIN(IFName, true, false);
1143bdd1243dSDimitry Andric   if (std::error_code EC = BufOrErr.getError())
1144bdd1243dSDimitry Andric     return createFileError(BundlerConfig.InputFileNames.front(), EC);
1145bdd1243dSDimitry Andric 
1146bdd1243dSDimitry Andric   ArchiveBuffers.push_back(std::move(*BufOrErr));
1147bdd1243dSDimitry Andric   Expected<std::unique_ptr<llvm::object::Archive>> LibOrErr =
1148bdd1243dSDimitry Andric       Archive::create(ArchiveBuffers.back()->getMemBufferRef());
1149bdd1243dSDimitry Andric   if (!LibOrErr)
1150bdd1243dSDimitry Andric     return LibOrErr.takeError();
1151bdd1243dSDimitry Andric 
1152bdd1243dSDimitry Andric   auto Archive = std::move(*LibOrErr);
1153bdd1243dSDimitry Andric 
1154bdd1243dSDimitry Andric   Error ArchiveErr = Error::success();
1155bdd1243dSDimitry Andric   auto ChildEnd = Archive->child_end();
1156bdd1243dSDimitry Andric 
1157bdd1243dSDimitry Andric   /// Iterate over all bundled code object files in the input archive.
1158bdd1243dSDimitry Andric   for (auto ArchiveIter = Archive->child_begin(ArchiveErr);
1159bdd1243dSDimitry Andric        ArchiveIter != ChildEnd; ++ArchiveIter) {
1160bdd1243dSDimitry Andric     if (ArchiveErr)
1161bdd1243dSDimitry Andric       return ArchiveErr;
1162bdd1243dSDimitry Andric     auto ArchiveChildNameOrErr = (*ArchiveIter).getName();
1163bdd1243dSDimitry Andric     if (!ArchiveChildNameOrErr)
1164bdd1243dSDimitry Andric       return ArchiveChildNameOrErr.takeError();
1165bdd1243dSDimitry Andric 
1166bdd1243dSDimitry Andric     StringRef BundledObjectFile = sys::path::filename(*ArchiveChildNameOrErr);
1167bdd1243dSDimitry Andric 
1168bdd1243dSDimitry Andric     auto CodeObjectBufferRefOrErr = (*ArchiveIter).getMemoryBufferRef();
1169bdd1243dSDimitry Andric     if (!CodeObjectBufferRefOrErr)
1170bdd1243dSDimitry Andric       return CodeObjectBufferRefOrErr.takeError();
1171bdd1243dSDimitry Andric 
1172bdd1243dSDimitry Andric     auto CodeObjectBuffer =
1173bdd1243dSDimitry Andric         MemoryBuffer::getMemBuffer(*CodeObjectBufferRefOrErr, false);
1174bdd1243dSDimitry Andric 
1175bdd1243dSDimitry Andric     Expected<std::unique_ptr<FileHandler>> FileHandlerOrErr =
1176bdd1243dSDimitry Andric         CreateFileHandler(*CodeObjectBuffer, BundlerConfig);
1177bdd1243dSDimitry Andric     if (!FileHandlerOrErr)
1178bdd1243dSDimitry Andric       return FileHandlerOrErr.takeError();
1179bdd1243dSDimitry Andric 
1180bdd1243dSDimitry Andric     std::unique_ptr<FileHandler> &FileHandler = *FileHandlerOrErr;
1181bdd1243dSDimitry Andric     assert(FileHandler &&
1182bdd1243dSDimitry Andric            "FileHandle creation failed for file in the archive!");
1183bdd1243dSDimitry Andric 
1184bdd1243dSDimitry Andric     if (Error ReadErr = FileHandler->ReadHeader(*CodeObjectBuffer))
1185bdd1243dSDimitry Andric       return ReadErr;
1186bdd1243dSDimitry Andric 
1187bdd1243dSDimitry Andric     Expected<std::optional<StringRef>> CurBundleIDOrErr =
1188bdd1243dSDimitry Andric         FileHandler->ReadBundleStart(*CodeObjectBuffer);
1189bdd1243dSDimitry Andric     if (!CurBundleIDOrErr)
1190bdd1243dSDimitry Andric       return CurBundleIDOrErr.takeError();
1191bdd1243dSDimitry Andric 
1192bdd1243dSDimitry Andric     std::optional<StringRef> OptionalCurBundleID = *CurBundleIDOrErr;
1193bdd1243dSDimitry Andric     // No device code in this child, skip.
1194bdd1243dSDimitry Andric     if (!OptionalCurBundleID)
1195bdd1243dSDimitry Andric       continue;
1196bdd1243dSDimitry Andric     StringRef CodeObject = *OptionalCurBundleID;
1197bdd1243dSDimitry Andric 
1198bdd1243dSDimitry Andric     // Process all bundle entries (CodeObjects) found in this child of input
1199bdd1243dSDimitry Andric     // archive.
1200bdd1243dSDimitry Andric     while (!CodeObject.empty()) {
1201bdd1243dSDimitry Andric       SmallVector<StringRef> CompatibleTargets;
1202bdd1243dSDimitry Andric       auto CodeObjectInfo = OffloadTargetInfo(CodeObject, BundlerConfig);
1203bdd1243dSDimitry Andric       if (CodeObjectInfo.hasHostKind()) {
1204bdd1243dSDimitry Andric         // Do nothing, we don't extract host code yet.
1205bdd1243dSDimitry Andric       } else if (getCompatibleOffloadTargets(CodeObjectInfo, CompatibleTargets,
1206bdd1243dSDimitry Andric                                              BundlerConfig)) {
1207bdd1243dSDimitry Andric         std::string BundleData;
1208bdd1243dSDimitry Andric         raw_string_ostream DataStream(BundleData);
1209bdd1243dSDimitry Andric         if (Error Err = FileHandler->ReadBundle(DataStream, *CodeObjectBuffer))
1210bdd1243dSDimitry Andric           return Err;
1211bdd1243dSDimitry Andric 
1212bdd1243dSDimitry Andric         for (auto &CompatibleTarget : CompatibleTargets) {
1213bdd1243dSDimitry Andric           SmallString<128> BundledObjectFileName;
1214bdd1243dSDimitry Andric           BundledObjectFileName.assign(BundledObjectFile);
1215bdd1243dSDimitry Andric           auto OutputBundleName =
1216bdd1243dSDimitry Andric               Twine(llvm::sys::path::stem(BundledObjectFileName) + "-" +
1217bdd1243dSDimitry Andric                     CodeObject +
1218bdd1243dSDimitry Andric                     getDeviceLibraryFileName(BundledObjectFileName,
1219bdd1243dSDimitry Andric                                              CodeObjectInfo.TargetID))
1220bdd1243dSDimitry Andric                   .str();
1221bdd1243dSDimitry Andric           // Replace ':' in optional target feature list with '_' to ensure
1222bdd1243dSDimitry Andric           // cross-platform validity.
1223bdd1243dSDimitry Andric           std::replace(OutputBundleName.begin(), OutputBundleName.end(), ':',
1224bdd1243dSDimitry Andric                        '_');
1225bdd1243dSDimitry Andric 
1226bdd1243dSDimitry Andric           std::unique_ptr<MemoryBuffer> MemBuf = MemoryBuffer::getMemBufferCopy(
1227bdd1243dSDimitry Andric               DataStream.str(), OutputBundleName);
1228bdd1243dSDimitry Andric           ArchiveBuffers.push_back(std::move(MemBuf));
1229bdd1243dSDimitry Andric           llvm::MemoryBufferRef MemBufRef =
1230bdd1243dSDimitry Andric               MemoryBufferRef(*(ArchiveBuffers.back()));
1231bdd1243dSDimitry Andric 
1232bdd1243dSDimitry Andric           // For inserting <CompatibleTarget, list<CodeObject>> entry in
1233bdd1243dSDimitry Andric           // OutputArchivesMap.
1234*06c3fb27SDimitry Andric           if (!OutputArchivesMap.contains(CompatibleTarget)) {
1235bdd1243dSDimitry Andric 
1236bdd1243dSDimitry Andric             std::vector<NewArchiveMember> ArchiveMembers;
1237bdd1243dSDimitry Andric             ArchiveMembers.push_back(NewArchiveMember(MemBufRef));
1238bdd1243dSDimitry Andric             OutputArchivesMap.insert_or_assign(CompatibleTarget,
1239bdd1243dSDimitry Andric                                                std::move(ArchiveMembers));
1240bdd1243dSDimitry Andric           } else {
1241bdd1243dSDimitry Andric             OutputArchivesMap[CompatibleTarget].push_back(
1242bdd1243dSDimitry Andric                 NewArchiveMember(MemBufRef));
1243bdd1243dSDimitry Andric           }
1244bdd1243dSDimitry Andric         }
1245bdd1243dSDimitry Andric       }
1246bdd1243dSDimitry Andric 
1247bdd1243dSDimitry Andric       if (Error Err = FileHandler->ReadBundleEnd(*CodeObjectBuffer))
1248bdd1243dSDimitry Andric         return Err;
1249bdd1243dSDimitry Andric 
1250bdd1243dSDimitry Andric       Expected<std::optional<StringRef>> NextTripleOrErr =
1251bdd1243dSDimitry Andric           FileHandler->ReadBundleStart(*CodeObjectBuffer);
1252bdd1243dSDimitry Andric       if (!NextTripleOrErr)
1253bdd1243dSDimitry Andric         return NextTripleOrErr.takeError();
1254bdd1243dSDimitry Andric 
1255bdd1243dSDimitry Andric       CodeObject = ((*NextTripleOrErr).has_value()) ? **NextTripleOrErr : "";
1256bdd1243dSDimitry Andric     } // End of processing of all bundle entries of this child of input archive.
1257bdd1243dSDimitry Andric   }   // End of while over children of input archive.
1258bdd1243dSDimitry Andric 
1259bdd1243dSDimitry Andric   assert(!ArchiveErr && "Error occurred while reading archive!");
1260bdd1243dSDimitry Andric 
1261bdd1243dSDimitry Andric   /// Write out an archive for each target
1262bdd1243dSDimitry Andric   for (auto &Target : BundlerConfig.TargetNames) {
1263bdd1243dSDimitry Andric     StringRef FileName = TargetOutputFileNameMap[Target];
1264bdd1243dSDimitry Andric     StringMapIterator<std::vector<llvm::NewArchiveMember>> CurArchiveMembers =
1265bdd1243dSDimitry Andric         OutputArchivesMap.find(Target);
1266bdd1243dSDimitry Andric     if (CurArchiveMembers != OutputArchivesMap.end()) {
1267bdd1243dSDimitry Andric       if (Error WriteErr = writeArchive(FileName, CurArchiveMembers->getValue(),
1268bdd1243dSDimitry Andric                                         true, getDefaultArchiveKindForHost(),
1269bdd1243dSDimitry Andric                                         true, false, nullptr))
1270bdd1243dSDimitry Andric         return WriteErr;
1271bdd1243dSDimitry Andric     } else if (!BundlerConfig.AllowMissingBundles) {
1272bdd1243dSDimitry Andric       std::string ErrMsg =
1273bdd1243dSDimitry Andric           Twine("no compatible code object found for the target '" + Target +
1274bdd1243dSDimitry Andric                 "' in heterogeneous archive library: " + IFName)
1275bdd1243dSDimitry Andric               .str();
1276bdd1243dSDimitry Andric       return createStringError(inconvertibleErrorCode(), ErrMsg);
1277bdd1243dSDimitry Andric     } else { // Create an empty archive file if no compatible code object is
1278bdd1243dSDimitry Andric              // found and "allow-missing-bundles" is enabled. It ensures that
1279bdd1243dSDimitry Andric              // the linker using output of this step doesn't complain about
1280bdd1243dSDimitry Andric              // the missing input file.
1281bdd1243dSDimitry Andric       std::vector<llvm::NewArchiveMember> EmptyArchive;
1282bdd1243dSDimitry Andric       EmptyArchive.clear();
1283bdd1243dSDimitry Andric       if (Error WriteErr = writeArchive(FileName, EmptyArchive, true,
1284bdd1243dSDimitry Andric                                         getDefaultArchiveKindForHost(), true,
1285bdd1243dSDimitry Andric                                         false, nullptr))
1286bdd1243dSDimitry Andric         return WriteErr;
1287bdd1243dSDimitry Andric     }
1288bdd1243dSDimitry Andric   }
1289bdd1243dSDimitry Andric 
1290bdd1243dSDimitry Andric   return Error::success();
1291bdd1243dSDimitry Andric }
1292