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