clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

   1 //===-- clang-offload-packager/ClangOffloadPackager.cpp - file bundler ---===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===---------------------------------------------------------------------===//
   8 //
   9 // This tool takes several device object files and bundles them into a single
  10 // binary image using a custom binary format. This is intended to be used to
  11 // embed many device files into an application to create a fat binary.
  12 //
  13 //===---------------------------------------------------------------------===//
  14
  15 #include "clang/Basic/Version.h"
  16
  17 #include "llvm/ADT/StringExtras.h"
  18 #include "llvm/BinaryFormat/Magic.h"
  19 #include "llvm/Object/ArchiveWriter.h"
  20 #include "llvm/Object/OffloadBinary.h"
  21 #include "llvm/Support/CommandLine.h"
  22 #include "llvm/Support/FileOutputBuffer.h"
  23 #include "llvm/Support/FileSystem.h"
  24 #include "llvm/Support/MemoryBuffer.h"
  25 #include "llvm/Support/Path.h"
  26 #include "llvm/Support/Signals.h"
  27 #include "llvm/Support/StringSaver.h"
  28 #include "llvm/Support/WithColor.h"
  29
  30 using namespace llvm;
  31 using namespace llvm::object;
  32
  33 static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
  34
  35 static cl::OptionCategory
  36     ClangOffloadPackagerCategory("clang-offload-packager options");
  37
  38 static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),
  39                                        cl::value_desc("file"),
  40                                        cl::cat(ClangOffloadPackagerCategory));
  41
  42 static cl::opt<std::string> InputFile(cl::Positional,
  43                                       cl::desc("Extract from <file>."),
  44                                       cl::value_desc("file"),
  45                                       cl::cat(ClangOffloadPackagerCategory));
  46
  47 static cl::list<std::string>
  48     DeviceImages("image",
  49                  cl::desc("List of key and value arguments. Required keywords "
  50                           "are 'file' and 'triple'."),
  51                  cl::value_desc("<key>=<value>,..."),
  52                  cl::cat(ClangOffloadPackagerCategory));
  53
  54 static cl::opt<bool>
  55     CreateArchive("archive",
  56                   cl::desc("Write extracted files to a static archive"),
  57                   cl::cat(ClangOffloadPackagerCategory));
  58
  59 /// Path of the current binary.
  60 static const char *PackagerExecutable;
  61
  62 static void PrintVersion(raw_ostream &OS) {
  63   OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n';
  64 }
  65
  66 // Get a map containing all the arguments for the image. Repeated arguments will
  67 // be placed in a comma separated list.
  68 static DenseMap<StringRef, StringRef> getImageArguments(StringRef Image,
  69                                                         StringSaver &Saver) {
  70   DenseMap<StringRef, StringRef> Args;
  71   for (StringRef Arg : llvm::split(Image, ",")) {
  72     auto [Key, Value] = Arg.split("=");
  73     if (Args.count(Key))
  74       Args[Key] = Saver.save(Args[Key] + "," + Value);
  75     else
  76       Args[Key] = Value;
  77   }
  78
  79   return Args;
  80 }
  81
  82 static Error writeFile(StringRef Filename, StringRef Data) {
  83   Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
  84       FileOutputBuffer::create(Filename, Data.size());
  85   if (!OutputOrErr)
  86     return OutputOrErr.takeError();
  87   std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
  88   llvm::copy(Data, Output->getBufferStart());
  89   if (Error E = Output->commit())
  90     return E;
  91   return Error::success();
  92 }
  93
  94 static Error bundleImages() {
  95   SmallVector<char, 1024> BinaryData;
  96   raw_svector_ostream OS(BinaryData);
  97   for (StringRef Image : DeviceImages) {
  98     BumpPtrAllocator Alloc;
  99     StringSaver Saver(Alloc);
 100     DenseMap<StringRef, StringRef> Args = getImageArguments(Image, Saver);
 101
 102     if (!Args.count("triple") || !Args.count("file"))
 103       return createStringError(
 104           inconvertibleErrorCode(),
 105           "'file' and 'triple' are required image arguments");
 106
 107     OffloadBinary::OffloadingImage ImageBinary{};
 108     std::unique_ptr<llvm::MemoryBuffer> DeviceImage;
 109     for (const auto &[Key, Value] : Args) {
 110       if (Key == "file") {
 111         llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ObjectOrErr =
 112             llvm::MemoryBuffer::getFileOrSTDIN(Value);
 113         if (std::error_code EC = ObjectOrErr.getError())
 114           return errorCodeToError(EC);
 115
 116         // Clang uses the '.o' suffix for LTO bitcode.
 117         if (identify_magic((*ObjectOrErr)->getBuffer()) == file_magic::bitcode)
 118           ImageBinary.TheImageKind = object::IMG_Bitcode;
 119         else
 120           ImageBinary.TheImageKind =
 121               getImageKind(sys::path::extension(Value).drop_front());
 122         ImageBinary.Image = std::move(*ObjectOrErr);
 123       } else if (Key == "kind") {
 124         ImageBinary.TheOffloadKind = getOffloadKind(Value);
 125       } else {
 126         ImageBinary.StringData[Key] = Value;
 127       }
 128     }
 129     llvm::SmallString<0> Buffer = OffloadBinary::write(ImageBinary);
 130     if (Buffer.size() % OffloadBinary::getAlignment() != 0)
 131       return createStringError(inconvertibleErrorCode(),
 132                                "Offload binary has invalid size alignment");
 133     OS << Buffer;
 134   }
 135
 136   if (Error E = writeFile(OutputFile,
 137                           StringRef(BinaryData.begin(), BinaryData.size())))
 138     return E;
 139   return Error::success();
 140 }
 141
 142 static Error unbundleImages() {
 143   ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
 144       MemoryBuffer::getFileOrSTDIN(InputFile);
 145   if (std::error_code EC = BufferOrErr.getError())
 146     return createFileError(InputFile, EC);
 147   std::unique_ptr<MemoryBuffer> Buffer = std::move(*BufferOrErr);
 148
 149   // This data can be misaligned if extracted from an archive.
 150   if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
 151                      Buffer->getBufferStart()))
 152     Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(),
 153                                             Buffer->getBufferIdentifier());
 154
 155   SmallVector<OffloadFile> Binaries;
 156   if (Error Err = extractOffloadBinaries(*Buffer, Binaries))
 157     return Err;
 158
 159   // Try to extract each device image specified by the user from the input file.
 160   for (StringRef Image : DeviceImages) {
 161     BumpPtrAllocator Alloc;
 162     StringSaver Saver(Alloc);
 163     auto Args = getImageArguments(Image, Saver);
 164
 165     SmallVector<const OffloadBinary *> Extracted;
 166     for (const OffloadFile &File : Binaries) {
 167       const auto *Binary = File.getBinary();
 168       // We handle the 'file' and 'kind' identifiers differently.
 169       bool Match = llvm::all_of(Args, [&](auto &Arg) {
 170         const auto [Key, Value] = Arg;
 171         if (Key == "file")
 172           return true;
 173         if (Key == "kind")
 174           return Binary->getOffloadKind() == getOffloadKind(Value);
 175         return Binary->getString(Key) == Value;
 176       });
 177       if (Match)
 178         Extracted.push_back(Binary);
 179     }
 180
 181     if (Extracted.empty())
 182       continue;
 183
 184     if (CreateArchive) {
 185       if (!Args.count("file"))
 186         return createStringError(inconvertibleErrorCode(),
 187                                  "Image must have a 'file' argument.");
 188
 189       SmallVector<NewArchiveMember> Members;
 190       for (const OffloadBinary *Binary : Extracted)
 191         Members.emplace_back(MemoryBufferRef(
 192             Binary->getImage(),
 193             Binary->getMemoryBufferRef().getBufferIdentifier()));
 194
 195       if (Error E = writeArchive(
 196               Args["file"], Members, SymtabWritingMode::NormalSymtab,
 197               Archive::getDefaultKindForHost(), true, false, nullptr))
 198         return E;
 199     } else if (Args.count("file")) {
 200       if (Extracted.size() > 1)
 201         WithColor::warning(errs(), PackagerExecutable)
 202             << "Multiple inputs match to a single file, '" << Args["file"]
 203             << "'\n";
 204       if (Error E = writeFile(Args["file"], Extracted.back()->getImage()))
 205         return E;
 206     } else {
 207       uint64_t Idx = 0;
 208       for (const OffloadBinary *Binary : Extracted) {
 209         StringRef Filename =
 210             Saver.save(sys::path::stem(InputFile) + "-" + Binary->getTriple() +
 211                        "-" + Binary->getArch() + "." + std::to_string(Idx++) +
 212                        "." + getImageKindName(Binary->getImageKind()));
 213         if (Error E = writeFile(Filename, Binary->getImage()))
 214           return E;
 215       }
 216     }
 217   }
 218
 219   return Error::success();
 220 }
 221
 222 int main(int argc, const char **argv) {
 223   sys::PrintStackTraceOnErrorSignal(argv[0]);
 224   cl::HideUnrelatedOptions(ClangOffloadPackagerCategory);
 225   cl::SetVersionPrinter(PrintVersion);
 226   cl::ParseCommandLineOptions(
 227       argc, argv,
 228       "A utility for bundling several object files into a single binary.\n"
 229       "The output binary can then be embedded into the host section table\n"
 230       "to create a fatbinary containing offloading code.\n");
 231
 232   if (Help) {
 233     cl::PrintHelpMessage();
 234     return EXIT_SUCCESS;
 235   }
 236
 237   PackagerExecutable = argv[0];
 238   auto reportError = [argv](Error E) {
 239     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
 240     return EXIT_FAILURE;
 241   };
 242
 243   if (!InputFile.empty() && !OutputFile.empty())
 244     return reportError(
 245         createStringError(inconvertibleErrorCode(),
 246                           "Packaging to an output file and extracting from an "
 247                           "input file are mutually exclusive."));
 248
 249   if (!OutputFile.empty()) {
 250     if (Error Err = bundleImages())
 251       return reportError(std::move(Err));
 252   } else if (!InputFile.empty()) {
 253     if (Error Err = unbundleImages())
 254       return reportError(std::move(Err));
 255   }
 256
 257   return EXIT_SUCCESS;
 258 }