259 lines · cpp
1//===-- llvm-offload-binary.cpp - offload binary management utility -------===//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-exception6//7//===----------------------------------------------------------------------===//8//9// This tool takes several device object files and bundles them into a single10// binary image using a custom binary format. This is intended to be used to11// embed many device files into an application to create a fat binary. It also12// supports extracting these files from a known location.13//14//===----------------------------------------------------------------------===//15 16#include "llvm/ADT/StringExtras.h"17#include "llvm/BinaryFormat/Magic.h"18#include "llvm/Object/ArchiveWriter.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 29using namespace llvm;30using namespace llvm::object;31 32static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);33 34static cl::OptionCategory OffloadBinaryCategory("llvm-offload-binary options");35 36static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),37 cl::value_desc("file"),38 cl::cat(OffloadBinaryCategory));39 40static cl::opt<std::string> InputFile(cl::Positional,41 cl::desc("Extract from <file>."),42 cl::value_desc("file"),43 cl::cat(OffloadBinaryCategory));44 45static 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(OffloadBinaryCategory));51 52static cl::opt<bool>53 CreateArchive("archive",54 cl::desc("Write extracted files to a static archive"),55 cl::cat(OffloadBinaryCategory));56 57/// Path of the current binary.58static const char *PackagerExecutable;59 60// Get a map containing all the arguments for the image. Repeated arguments will61// be placed in a comma separated list.62static DenseMap<StringRef, StringRef> getImageArguments(StringRef Image,63 StringSaver &Saver) {64 DenseMap<StringRef, StringRef> Args;65 for (StringRef Arg : llvm::split(Image, ",")) {66 auto [Key, Value] = Arg.split("=");67 auto [It, Inserted] = Args.try_emplace(Key, Value);68 if (!Inserted)69 It->second = Saver.save(It->second + "," + Value);70 }71 72 return Args;73}74 75static Error writeFile(StringRef Filename, StringRef Data) {76 Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =77 FileOutputBuffer::create(Filename, Data.size());78 if (!OutputOrErr)79 return OutputOrErr.takeError();80 std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);81 llvm::copy(Data, Output->getBufferStart());82 if (Error E = Output->commit())83 return E;84 return Error::success();85}86 87static Error bundleImages() {88 SmallVector<char, 1024> BinaryData;89 raw_svector_ostream OS(BinaryData);90 for (StringRef Image : DeviceImages) {91 BumpPtrAllocator Alloc;92 StringSaver Saver(Alloc);93 DenseMap<StringRef, StringRef> Args = getImageArguments(Image, Saver);94 95 if (!Args.count("file"))96 return createStringError(inconvertibleErrorCode(),97 "'file' is a required image arguments");98 99 // Permit using multiple instances of `file` in a single string.100 for (auto &File : llvm::split(Args["file"], ",")) {101 OffloadBinary::OffloadingImage ImageBinary{};102 103 llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ObjectOrErr =104 llvm::MemoryBuffer::getFileOrSTDIN(File);105 if (std::error_code EC = ObjectOrErr.getError())106 return errorCodeToError(EC);107 108 // Clang uses the '.o' suffix for LTO bitcode.109 if (identify_magic((*ObjectOrErr)->getBuffer()) == file_magic::bitcode)110 ImageBinary.TheImageKind = object::IMG_Bitcode;111 else if (sys::path::has_extension(File))112 ImageBinary.TheImageKind =113 getImageKind(sys::path::extension(File).drop_front());114 else115 ImageBinary.TheImageKind = IMG_None;116 ImageBinary.Image = std::move(*ObjectOrErr);117 for (const auto &[Key, Value] : Args) {118 if (Key == "kind") {119 ImageBinary.TheOffloadKind = getOffloadKind(Value);120 } else if (Key != "file") {121 ImageBinary.StringData[Key] = Value;122 }123 }124 llvm::SmallString<0> Buffer = OffloadBinary::write(ImageBinary);125 if (Buffer.size() % OffloadBinary::getAlignment() != 0)126 return createStringError(inconvertibleErrorCode(),127 "Offload binary has invalid size alignment");128 OS << Buffer;129 }130 }131 132 if (Error E = writeFile(OutputFile,133 StringRef(BinaryData.begin(), BinaryData.size())))134 return E;135 return Error::success();136}137 138static Error unbundleImages() {139 ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =140 MemoryBuffer::getFileOrSTDIN(InputFile);141 if (std::error_code EC = BufferOrErr.getError())142 return createFileError(InputFile, EC);143 std::unique_ptr<MemoryBuffer> Buffer = std::move(*BufferOrErr);144 145 // This data can be misaligned if extracted from an archive.146 if (!isAddrAligned(Align(OffloadBinary::getAlignment()),147 Buffer->getBufferStart()))148 Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(),149 Buffer->getBufferIdentifier());150 151 SmallVector<OffloadFile> Binaries;152 if (Error Err = extractOffloadBinaries(*Buffer, Binaries))153 return Err;154 155 // Try to extract each device image specified by the user from the input file.156 for (StringRef Image : DeviceImages) {157 BumpPtrAllocator Alloc;158 StringSaver Saver(Alloc);159 auto Args = getImageArguments(Image, Saver);160 161 SmallVector<const OffloadBinary *> Extracted;162 for (const OffloadFile &File : Binaries) {163 const auto *Binary = File.getBinary();164 // We handle the 'file' and 'kind' identifiers differently.165 bool Match = llvm::all_of(Args, [&](auto &Arg) {166 const auto [Key, Value] = Arg;167 if (Key == "file")168 return true;169 if (Key == "kind")170 return Binary->getOffloadKind() == getOffloadKind(Value);171 return Binary->getString(Key) == Value;172 });173 if (Match)174 Extracted.push_back(Binary);175 }176 177 if (Extracted.empty())178 continue;179 180 if (CreateArchive) {181 if (!Args.count("file"))182 return createStringError(inconvertibleErrorCode(),183 "Image must have a 'file' argument.");184 185 SmallVector<NewArchiveMember> Members;186 for (const OffloadBinary *Binary : Extracted)187 Members.emplace_back(MemoryBufferRef(188 Binary->getImage(),189 Binary->getMemoryBufferRef().getBufferIdentifier()));190 191 if (Error E = writeArchive(192 Args["file"], Members, SymtabWritingMode::NormalSymtab,193 Archive::getDefaultKind(), true, false, nullptr))194 return E;195 } else if (auto It = Args.find("file"); It != Args.end()) {196 if (Extracted.size() > 1)197 WithColor::warning(errs(), PackagerExecutable)198 << "Multiple inputs match to a single file, '" << It->second199 << "'\n";200 if (Error E = writeFile(It->second, Extracted.back()->getImage()))201 return E;202 } else {203 uint64_t Idx = 0;204 for (const OffloadBinary *Binary : Extracted) {205 StringRef Filename =206 Saver.save(sys::path::stem(InputFile) + "-" + Binary->getTriple() +207 "-" + Binary->getArch() + "." + std::to_string(Idx++) +208 "." + getImageKindName(Binary->getImageKind()));209 if (Error E = writeFile(Filename, Binary->getImage()))210 return E;211 }212 }213 }214 215 return Error::success();216}217 218int main(int argc, const char **argv) {219 sys::PrintStackTraceOnErrorSignal(argv[0]);220 cl::HideUnrelatedOptions(OffloadBinaryCategory);221 cl::ParseCommandLineOptions(222 argc, argv,223 "A utility for bundling several object files into a single binary.\n"224 "The output binary can then be embedded into the host section table\n"225 "to create a fatbinary containing offloading code.\n");226 227 if (sys::path::stem(argv[0]).ends_with("clang-offload-packager"))228 WithColor::warning(errs(), PackagerExecutable)229 << "'clang-offload-packager' is deprecated. Use 'llvm-offload-binary' "230 "instead.\n";231 232 if (Help || (OutputFile.empty() && InputFile.empty())) {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}259