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 clEnumValN(object::OFK_SYCL, "sycl", "Wrap SYCL binaries")));
45
46static cl::opt<bool> Relocatable(
47 "relocatable",
48 cl::desc("Wrap for a relocatable offloading application (OpenMP only)"),
49 cl::cat(OffloadWrapeprCategory));
50
51static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),
52 cl::value_desc("file"),
53 cl::cat(OffloadWrapeprCategory));
54
55static cl::list<std::string> InputFiles(cl::Positional,
56 cl::desc("Wrap input from <file>"),
57 cl::value_desc("file"), cl::OneOrMore,
58 cl::cat(OffloadWrapeprCategory));
59
60static cl::opt<std::string>
61 TheTriple("triple", cl::desc("Target triple for the wrapper module"),
62 cl::init(Val: sys::getDefaultTargetTriple()),
63 cl::cat(OffloadWrapeprCategory));
64
65static Error wrapImages(ArrayRef<ArrayRef<char>> BuffersToWrap) {
66 if (BuffersToWrap.size() > 1 &&
67 (Kind == llvm::object::OFK_Cuda || Kind == llvm::object::OFK_HIP))
68 return createStringError(
69 Fmt: "CUDA / HIP offloading uses a single fatbinary or offload bundle");
70
71 LLVMContext Context;
72 Module M("offload.wrapper.module", Context);
73 M.setTargetTriple(llvm::Triple(TheTriple));
74
75 switch (Kind) {
76 case llvm::object::OFK_OpenMP:
77 if (Error Err = offloading::wrapOpenMPBinaries(
78 M, Images: BuffersToWrap, EntryArray: offloading::getOffloadEntryArray(M),
79 /*Suffix=*/"", /*Relocatable=*/Relocatable))
80 return Err;
81 break;
82 case llvm::object::OFK_Cuda:
83 if (Error Err = offloading::wrapCudaBinary(
84 M, Images: BuffersToWrap.front(), EntryArray: offloading::getOffloadEntryArray(M),
85 /*Suffix=*/"", /*EmitSurfacesAndTextures=*/false))
86 return Err;
87 break;
88 case llvm::object::OFK_HIP:
89 if (Error Err = offloading::wrapHIPBinary(
90 M, Images: BuffersToWrap.front(), EntryArray: offloading::getOffloadEntryArray(M)))
91 return Err;
92 break;
93 case llvm::object::OFK_SYCL:
94 if (Error Err = offloading::wrapSYCLBinaries(M, Buffer: BuffersToWrap.front()))
95 return Err;
96 break;
97 default:
98 return createStringError(S: getOffloadKindName(Name: Kind) +
99 " wrapping is not supported");
100 }
101
102 int FD = -1;
103 if (std::error_code EC = sys::fs::openFileForWrite(Name: OutputFile, ResultFD&: FD))
104 return errorCodeToError(EC);
105 llvm::raw_fd_ostream OS(FD, true);
106 WriteBitcodeToFile(M, Out&: OS);
107
108 return Error::success();
109}
110
111int main(int argc, char **argv) {
112 InitLLVM X(argc, argv);
113 cl::HideUnrelatedOptions(Category&: OffloadWrapeprCategory);
114 cl::ParseCommandLineOptions(
115 argc, argv,
116 Overview: "Generate runtime registration code for a device binary image\n");
117
118 if (Help) {
119 cl::PrintHelpMessage();
120 return EXIT_SUCCESS;
121 }
122
123 auto ReportError = [argv](Error E) {
124 logAllUnhandledErrors(E: std::move(E), OS&: WithColor::error(OS&: errs(), Prefix: argv[0]));
125 exit(EXIT_FAILURE);
126 };
127
128 SmallVector<std::unique_ptr<MemoryBuffer>> Buffers;
129 SmallVector<ArrayRef<char>> BuffersToWrap;
130 for (StringRef Input : InputFiles) {
131 ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
132 MemoryBuffer::getFileOrSTDIN(Filename: Input);
133 if (std::error_code EC = BufferOrErr.getError())
134 ReportError(createFileError(F: Input, EC));
135 std::unique_ptr<MemoryBuffer> &Buffer =
136 Buffers.emplace_back(Args: std::move(*BufferOrErr));
137 BuffersToWrap.emplace_back(
138 Args: ArrayRef<char>(Buffer->getBufferStart(), Buffer->getBufferSize()));
139 }
140
141 if (Error Err = wrapImages(BuffersToWrap))
142 ReportError(std::move(Err));
143
144 return EXIT_SUCCESS;
145}
146