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