140 lines · cpp
1//===- llvm-offload-wrapper: Create runtime registration code for devices -===//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// Provides a utility for generating runtime registration code for device code.10// We take a binary image (CUDA fatbinary, HIP offload bundle, LLVM binary) and11// create a new IR module that calls the respective runtime to load it on the12// device.13//14//===----------------------------------------------------------------------===//15 16#include "llvm/Bitcode/BitcodeWriter.h"17#include "llvm/Frontend/Offloading/OffloadWrapper.h"18#include "llvm/Frontend/Offloading/Utility.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/InitLLVM.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#include "llvm/TargetParser/Host.h"30 31using namespace llvm;32 33static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);34 35static cl::OptionCategory36 OffloadWrapeprCategory("llvm-offload-wrapper options");37 38static cl::opt<object::OffloadKind> Kind(39 "kind", cl::desc("Wrap for offload kind:"), cl::cat(OffloadWrapeprCategory),40 cl::Required,41 cl::values(clEnumValN(object::OFK_OpenMP, "openmp", "Wrap OpenMP binaries"),42 clEnumValN(object::OFK_Cuda, "cuda", "Wrap CUDA binaries"),43 clEnumValN(object::OFK_HIP, "hip", "Wrap HIP binaries")));44 45static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),46 cl::value_desc("file"),47 cl::cat(OffloadWrapeprCategory));48 49static cl::list<std::string> InputFiles(cl::Positional,50 cl::desc("Wrap input from <file>"),51 cl::value_desc("file"), cl::OneOrMore,52 cl::cat(OffloadWrapeprCategory));53 54static cl::opt<std::string>55 TheTriple("triple", cl::desc("Target triple for the wrapper module"),56 cl::init(sys::getDefaultTargetTriple()),57 cl::cat(OffloadWrapeprCategory));58 59static Error wrapImages(ArrayRef<ArrayRef<char>> BuffersToWrap) {60 if (BuffersToWrap.size() > 1 &&61 (Kind == llvm::object::OFK_Cuda || Kind == llvm::object::OFK_HIP))62 return createStringError(63 "CUDA / HIP offloading uses a single fatbinary or offload bundle");64 65 LLVMContext Context;66 Module M("offload.wrapper.module", Context);67 M.setTargetTriple(llvm::Triple(TheTriple));68 69 switch (Kind) {70 case llvm::object::OFK_OpenMP:71 if (Error Err = offloading::wrapOpenMPBinaries(72 M, BuffersToWrap, offloading::getOffloadEntryArray(M),73 /*Suffix=*/"", /*Relocatable=*/false))74 return Err;75 break;76 case llvm::object::OFK_Cuda:77 if (Error Err = offloading::wrapCudaBinary(78 M, BuffersToWrap.front(), offloading::getOffloadEntryArray(M),79 /*Suffix=*/"", /*EmitSurfacesAndTextures=*/false))80 return Err;81 break;82 case llvm::object::OFK_HIP:83 if (Error Err = offloading::wrapHIPBinary(84 M, BuffersToWrap.front(), offloading::getOffloadEntryArray(M)))85 return Err;86 break;87 case llvm::object::OFK_SYCL:88 if (Error Err = offloading::wrapSYCLBinaries(M, BuffersToWrap.front()))89 return Err;90 break;91 default:92 return createStringError(getOffloadKindName(Kind) +93 " wrapping is not supported");94 }95 96 int FD = -1;97 if (std::error_code EC = sys::fs::openFileForWrite(OutputFile, FD))98 return errorCodeToError(EC);99 llvm::raw_fd_ostream OS(FD, true);100 WriteBitcodeToFile(M, OS);101 102 return Error::success();103}104 105int main(int argc, char **argv) {106 InitLLVM X(argc, argv);107 cl::HideUnrelatedOptions(OffloadWrapeprCategory);108 cl::ParseCommandLineOptions(109 argc, argv,110 "Generate runtime registration code for a device binary image\n");111 112 if (Help) {113 cl::PrintHelpMessage();114 return EXIT_SUCCESS;115 }116 117 auto ReportError = [argv](Error E) {118 logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));119 exit(EXIT_FAILURE);120 };121 122 SmallVector<std::unique_ptr<MemoryBuffer>> Buffers;123 SmallVector<ArrayRef<char>> BuffersToWrap;124 for (StringRef Input : InputFiles) {125 ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =126 MemoryBuffer::getFileOrSTDIN(Input);127 if (std::error_code EC = BufferOrErr.getError())128 ReportError(createFileError(Input, EC));129 std::unique_ptr<MemoryBuffer> &Buffer =130 Buffers.emplace_back(std::move(*BufferOrErr));131 BuffersToWrap.emplace_back(132 ArrayRef<char>(Buffer->getBufferStart(), Buffer->getBufferSize()));133 }134 135 if (Error Err = wrapImages(BuffersToWrap))136 ReportError(std::move(Err));137 138 return EXIT_SUCCESS;139}140