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/BinaryFormat/Magic.h"
  18 #include "llvm/Object/OffloadBinary.h"
  19 #include "llvm/Support/CommandLine.h"
  20 #include "llvm/Support/FileOutputBuffer.h"
  21 #include "llvm/Support/FileSystem.h"
  22 #include "llvm/Support/MemoryBuffer.h"
  23 #include "llvm/Support/Path.h"
  24 #include "llvm/Support/Signals.h"
  25 #include "llvm/Support/StringSaver.h"
  26 #include "llvm/Support/WithColor.h"
  27
  28 using namespace llvm;
  29 using namespace llvm::object;
  30
  31 static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
  32
  33 static cl::OptionCategory
  34     ClangOffloadPackagerCategory("clang-offload-packager options");
  35
  36 static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),
  37                                        cl::value_desc("file"),
  38                                        cl::cat(ClangOffloadPackagerCategory));
  39
  40 static cl::opt<std::string> InputFile(cl::Positional,
  41                                       cl::desc("Extract from <file>."),
  42                                       cl::value_desc("file"),
  43                                       cl::cat(ClangOffloadPackagerCategory));
  44
  45 static cl::list<std::string>
  46     DeviceImages("image",
  47                  cl::desc("List of key and value arguments. Required keywords "
  48                           "are 'file' and 'triple'."),
  49                  cl::value_desc("<key>=<value>,..."),
  50                  cl::cat(ClangOffloadPackagerCategory));
  51
  52 static void PrintVersion(raw_ostream &OS) {
  53   OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n';
  54 }
  55
  56 // Get a map containing all the arguments for the image. Repeated arguments will
  57 // be placed in a comma separated list.
  58 static DenseMap<StringRef, StringRef> getImageArguments(StringRef Image,
  59                                                         StringSaver &Saver) {
  60   DenseMap<StringRef, StringRef> Args;
  61   for (StringRef Arg : llvm::split(Image, ",")) {
  62     auto [Key, Value] = Arg.split("=");
  63     if (Args.count(Key))
  64       Args[Key] = Saver.save(Args[Key] + "," + Value);
  65     else
  66       Args[Key] = Value;
  67   }
  68
  69   return Args;
  70 }
  71
  72 static Error bundleImages() {
  73   SmallVector<char, 1024> BinaryData;
  74   raw_svector_ostream OS(BinaryData);
  75   for (StringRef Image : DeviceImages) {
  76     BumpPtrAllocator Alloc;
  77     StringSaver Saver(Alloc);
  78     DenseMap<StringRef, StringRef> Args = getImageArguments(Image, Saver);
  79
  80     if (!Args.count("triple") || !Args.count("file"))
  81       return createStringError(
  82           inconvertibleErrorCode(),
  83           "'file' and 'triple' are required image arguments");
  84
  85     OffloadBinary::OffloadingImage ImageBinary{};
  86     std::unique_ptr<llvm::MemoryBuffer> DeviceImage;
  87     for (const auto &[Key, Value] : Args) {
  88       if (Key == "file") {
  89         llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ObjectOrErr =
  90             llvm::MemoryBuffer::getFileOrSTDIN(Value);
  91         if (std::error_code EC = ObjectOrErr.getError())
  92           return errorCodeToError(EC);
  93
  94         // Clang uses the '.o' suffix for LTO bitcode.
  95         if (identify_magic((*ObjectOrErr)->getBuffer()) == file_magic::bitcode)
  96           ImageBinary.TheImageKind = object::IMG_Bitcode;
  97         else
  98           ImageBinary.TheImageKind =
  99               getImageKind(sys::path::extension(Value).drop_front());
 100         ImageBinary.Image = std::move(*ObjectOrErr);
 101       } else if (Key == "kind") {
 102         ImageBinary.TheOffloadKind = getOffloadKind(Value);
 103       } else {
 104         ImageBinary.StringData[Key] = Value;
 105       }
 106     }
 107     std::unique_ptr<MemoryBuffer> Buffer = OffloadBinary::write(ImageBinary);
 108     if (Buffer->getBufferSize() % OffloadBinary::getAlignment() != 0)
 109       return createStringError(inconvertibleErrorCode(),
 110                                "Offload binary has invalid size alignment");
 111     OS << Buffer->getBuffer();
 112   }
 113
 114   Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
 115       FileOutputBuffer::create(OutputFile, BinaryData.size());
 116   if (!OutputOrErr)
 117     return OutputOrErr.takeError();
 118   std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
 119   std::copy(BinaryData.begin(), BinaryData.end(), Output->getBufferStart());
 120   if (Error E = Output->commit())
 121     return E;
 122   return Error::success();
 123 }
 124
 125 static Error unbundleImages() {
 126   ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
 127       MemoryBuffer::getFileOrSTDIN(InputFile);
 128   if (std::error_code EC = BufferOrErr.getError())
 129     return createFileError(InputFile, EC);
 130   std::unique_ptr<MemoryBuffer> Buffer = std::move(*BufferOrErr);
 131
 132   // This data can be misaligned if extracted from an archive.
 133   if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
 134                      Buffer->getBufferStart()))
 135     Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(),
 136                                             Buffer->getBufferIdentifier());
 137
 138   SmallVector<OffloadFile> Binaries;
 139   if (Error Err = extractOffloadBinaries(*Buffer, Binaries))
 140     return Err;
 141
 142   // Try to extract each device image specified by the user from the input file.
 143   for (StringRef Image : DeviceImages) {
 144     BumpPtrAllocator Alloc;
 145     StringSaver Saver(Alloc);
 146     auto Args = getImageArguments(Image, Saver);
 147
 148     for (uint64_t I = 0, E = Binaries.size(); I != E; ++I) {
 149       const auto *Binary = Binaries[I].getBinary();
 150       // We handle the 'file' and 'kind' identifiers differently.
 151       bool Match = llvm::all_of(Args, [&](auto &Arg) {
 152         const auto [Key, Value] = Arg;
 153         if (Key == "file")
 154           return true;
 155         if (Key == "kind")
 156           return Binary->getOffloadKind() == getOffloadKind(Value);
 157         return Binary->getString(Key) == Value;
 158       });
 159       if (!Match)
 160         continue;
 161
 162       // If the user did not provide a filename derive one from the input and
 163       // image.
 164       StringRef Filename =
 165           !Args.count("file")
 166               ? Saver.save(sys::path::stem(InputFile) + "-" +
 167                            Binary->getTriple() + "-" + Binary->getArch() + "." +
 168                            std::to_string(I) + "." +
 169                            getImageKindName(Binary->getImageKind()))
 170               : Args["file"];
 171
 172       Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
 173           FileOutputBuffer::create(Filename, Binary->getImage().size());
 174       if (!OutputOrErr)
 175         return OutputOrErr.takeError();
 176       std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
 177       llvm::copy(Binary->getImage(), Output->getBufferStart());
 178       if (Error E = Output->commit())
 179         return E;
 180     }
 181   }
 182
 183   return Error::success();
 184 }
 185
 186 int main(int argc, const char **argv) {
 187   sys::PrintStackTraceOnErrorSignal(argv[0]);
 188   cl::HideUnrelatedOptions(ClangOffloadPackagerCategory);
 189   cl::SetVersionPrinter(PrintVersion);
 190   cl::ParseCommandLineOptions(
 191       argc, argv,
 192       "A utility for bundling several object files into a single binary.\n"
 193       "The output binary can then be embedded into the host section table\n"
 194       "to create a fatbinary containing offloading code.\n");
 195
 196   if (Help) {
 197     cl::PrintHelpMessage();
 198     return EXIT_SUCCESS;
 199   }
 200
 201   auto reportError = [argv](Error E) {
 202     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
 203     return EXIT_FAILURE;
 204   };
 205
 206   if (!InputFile.empty() && !OutputFile.empty())
 207     return reportError(
 208         createStringError(inconvertibleErrorCode(),
 209                           "Packaging to an output file and extracting from an "
 210                           "input file are mutually exclusive."));
 211
 212   if (!OutputFile.empty()) {
 213     if (Error Err = bundleImages())
 214       return reportError(std::move(Err));
 215   } else if (!InputFile.empty()) {
 216     if (Error Err = unbundleImages())
 217       return reportError(std::move(Err));
 218   }
 219
 220   return EXIT_SUCCESS;
 221 }