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-exception
6//
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) and
11// create a new IR module that calls the respective runtime to load it on the
12// 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::OptionCategory
36 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(Val: 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 Fmt: "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, Images: BuffersToWrap, EntryArray: 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, Images: BuffersToWrap.front(), EntryArray: 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, Images: BuffersToWrap.front(), EntryArray: offloading::getOffloadEntryArray(M)))
85 return Err;
86 break;
87 case llvm::object::OFK_SYCL:
88 if (Error Err = offloading::wrapSYCLBinaries(M, Buffer: BuffersToWrap.front()))
89 return Err;
90 break;
91 default:
92 return createStringError(S: getOffloadKindName(Name: Kind) +
93 " wrapping is not supported");
94 }
95
96 int FD = -1;
97 if (std::error_code EC = sys::fs::openFileForWrite(Name: OutputFile, ResultFD&: FD))
98 return errorCodeToError(EC);
99 llvm::raw_fd_ostream OS(FD, true);
100 WriteBitcodeToFile(M, Out&: OS);
101
102 return Error::success();
103}
104
105int main(int argc, char **argv) {
106 InitLLVM X(argc, argv);
107 cl::HideUnrelatedOptions(Category&: OffloadWrapeprCategory);
108 cl::ParseCommandLineOptions(
109 argc, argv,
110 Overview: "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(E: std::move(E), OS&: WithColor::error(OS&: errs(), Prefix: 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(Filename: Input);
127 if (std::error_code EC = BufferOrErr.getError())
128 ReportError(createFileError(F: Input, EC));
129 std::unique_ptr<MemoryBuffer> &Buffer =
130 Buffers.emplace_back(Args: std::move(*BufferOrErr));
131 BuffersToWrap.emplace_back(
132 Args: 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