1//===-- Main entry into the loader interface ------------------------------===//
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// This utility is used to launch standard programs onto the GPU in conjunction
10// with the LLVM 'libc' project. It is designed to mimic a standard emulator
11// workflow, allowing for unit tests to be run on the GPU directly.
12//
13//===----------------------------------------------------------------------===//
14
15#include "llvm-gpu-loader.h"
16
17#include "llvm/BinaryFormat/Magic.h"
18#include "llvm/Object/ELF.h"
19#include "llvm/Object/ELFObjectFile.h"
20#include "llvm/Support/CommandLine.h"
21#include "llvm/Support/Error.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/WithColor.h"
27#include "llvm/TargetParser/Triple.h"
28
29#include <cerrno>
30#include <cstdio>
31#include <cstdlib>
32#include <cstring>
33#include <string>
34
35using namespace llvm;
36
37static cl::OptionCategory LoaderCategory("loader options");
38
39static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden,
40 cl::cat(LoaderCategory));
41
42static cl::opt<unsigned>
43 ThreadsX("threads-x", cl::desc("Number of threads in the 'x' dimension"),
44 cl::init(Val: 1), cl::cat(LoaderCategory));
45static cl::opt<unsigned>
46 ThreadsY("threads-y", cl::desc("Number of threads in the 'y' dimension"),
47 cl::init(Val: 1), cl::cat(LoaderCategory));
48static cl::opt<unsigned>
49 ThreadsZ("threads-z", cl::desc("Number of threads in the 'z' dimension"),
50 cl::init(Val: 1), cl::cat(LoaderCategory));
51static cl::alias threads("threads", cl::aliasopt(ThreadsX),
52 cl::desc("Alias for --threads-x"),
53 cl::cat(LoaderCategory));
54
55static cl::opt<unsigned>
56 BlocksX("blocks-x", cl::desc("Number of blocks in the 'x' dimension"),
57 cl::init(Val: 1), cl::cat(LoaderCategory));
58static cl::opt<unsigned>
59 BlocksY("blocks-y", cl::desc("Number of blocks in the 'y' dimension"),
60 cl::init(Val: 1), cl::cat(LoaderCategory));
61static cl::opt<unsigned>
62 BlocksZ("blocks-z", cl::desc("Number of blocks in the 'z' dimension"),
63 cl::init(Val: 1), cl::cat(LoaderCategory));
64static cl::alias Blocks("blocks", cl::aliasopt(BlocksX),
65 cl::desc("Alias for --blocks-x"),
66 cl::cat(LoaderCategory));
67
68static cl::list<std::string> Kernels(
69 "kernel", cl::value_desc("name"),
70 cl::desc("Launch '<name>(void)' instead of the 'main' entry point."),
71 cl::cat(LoaderCategory));
72
73static cl::opt<std::string> File(cl::Positional, cl::Required,
74 cl::desc("<gpu executable>"),
75 cl::cat(LoaderCategory));
76static cl::list<std::string> Args(cl::ConsumeAfter,
77 cl::desc("<program arguments>..."),
78 cl::cat(LoaderCategory));
79
80[[noreturn]] static void handleError(Error E) {
81 outs().flush();
82 logAllUnhandledErrors(E: std::move(E), OS&: WithColor::error(OS&: errs(), Prefix: "loader"));
83 exit(EXIT_FAILURE);
84}
85
86[[noreturn]] static void handleError(ol_result_t Err, unsigned Line) {
87 fprintf(stderr, format: "%s:%d %s\n", __FILE__, Line, Err->Details);
88 exit(EXIT_FAILURE);
89}
90
91#define OFFLOAD_ERR(X) \
92 if (ol_result_t Err = X) \
93 handleError(Err, __LINE__);
94
95static void *copyArgumentVector(int Argc, const char **Argv,
96 ol_device_handle_t Device) {
97 size_t ArgSize = sizeof(char *) * (Argc + 1);
98 size_t StringLen = 0;
99 for (int i = 0; i < Argc; ++i)
100 StringLen += strlen(s: Argv[i]) + 1;
101
102 // We allocate enough space for a null terminated array and all the strings.
103 void *DevArgv;
104 OFFLOAD_ERR(olMemAllocHost(Device, ArgSize + StringLen, &DevArgv));
105 if (!DevArgv)
106 handleError(
107 E: createStringError(Fmt: "Failed to allocate memory for environment."));
108
109 // Store the strings linerally in the same memory buffer.
110 void *DevString = reinterpret_cast<uint8_t *>(DevArgv) + ArgSize;
111 for (int i = 0; i < Argc; ++i) {
112 size_t size = strlen(s: Argv[i]) + 1;
113 std::memcpy(dest: DevString, src: Argv[i], n: size);
114 static_cast<void **>(DevArgv)[i] = DevString;
115 DevString = reinterpret_cast<uint8_t *>(DevString) + size;
116 }
117
118 // Ensure the vector is null terminated.
119 reinterpret_cast<void **>(DevArgv)[Argc] = nullptr;
120 return DevArgv;
121}
122
123void *copyEnvironment(const char **Envp, ol_device_handle_t Device) {
124 int Envc = 0;
125 for (const char **Env = Envp; *Env != 0; ++Env)
126 ++Envc;
127
128 return copyArgumentVector(Argc: Envc, Argv: Envp, Device);
129}
130
131ol_device_handle_t findDevice(MemoryBufferRef Binary) {
132 ol_device_handle_t Device = nullptr;
133 std::tuple Data = std::make_tuple(args: &Device, args: &Binary);
134 OFFLOAD_ERR(olIterateDevices(
135 [](ol_device_handle_t Device, void *UserData) {
136 auto &[Output, Binary] = *reinterpret_cast<decltype(Data) *>(UserData);
137 bool IsValid = false;
138 OFFLOAD_ERR(olIsValidBinary(Device, Binary->getBufferStart(),
139 Binary->getBufferSize(), &IsValid));
140 if (!IsValid)
141 return true;
142
143 *Output = Device;
144 return false;
145 },
146 &Data));
147 return Device;
148}
149
150ol_device_handle_t getHostDevice() {
151 ol_device_handle_t Device;
152 OFFLOAD_ERR(olIterateDevices(
153 [](ol_device_handle_t Device, void *UserData) {
154 ol_platform_handle_t Platform;
155 olGetDeviceInfo(Device, OL_DEVICE_INFO_PLATFORM, sizeof(Platform),
156 &Platform);
157 ol_platform_backend_t Backend;
158 olGetPlatformInfo(Platform, OL_PLATFORM_INFO_BACKEND, sizeof(Backend),
159 &Backend);
160
161 auto &Output = *reinterpret_cast<decltype(Device) *>(UserData);
162 if (Backend == OL_PLATFORM_BACKEND_HOST) {
163 Output = Device;
164 return false;
165 }
166 return true;
167 },
168 &Device));
169 return Device;
170}
171
172template <typename... Args>
173void launchKernel(ol_queue_handle_t Queue, ol_device_handle_t Device,
174 ol_program_handle_t Program, const char *Name,
175 ol_kernel_launch_size_args_t LaunchArgs,
176 Args &...KernelArgs) {
177 ol_symbol_handle_t Kernel;
178 OFFLOAD_ERR(olGetSymbol(Program, Name, OL_SYMBOL_KIND_KERNEL, &Kernel));
179
180 if constexpr (sizeof...(Args) == 0) {
181 OFFLOAD_ERR(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr, 0,
182 nullptr, nullptr));
183 } else {
184 void *ArgPtrs[] = {static_cast<void *>(&KernelArgs)...};
185 size_t ArgSizes[] = {sizeof(KernelArgs)...};
186 OFFLOAD_ERR(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
187 sizeof...(Args), ArgPtrs, ArgSizes));
188 }
189}
190
191int main(int argc, const char **argv, const char **envp) {
192 sys::PrintStackTraceOnErrorSignal(Argv0: argv[0]);
193 cl::HideUnrelatedOptions(Category&: LoaderCategory);
194 cl::ParseCommandLineOptions(
195 argc, argv,
196 Overview: "A utility used to launch unit tests built for a GPU target. This is\n"
197 "intended to provide an interface similar to cross-compiling "
198 "emulators\n");
199
200 if (Help) {
201 cl::PrintHelpMessage();
202 return EXIT_SUCCESS;
203 }
204
205 if (Error Err = loadLLVMOffload())
206 handleError(E: std::move(Err));
207
208 ErrorOr<std::unique_ptr<MemoryBuffer>> ImageOrErr =
209 MemoryBuffer::getFileOrSTDIN(Filename: File);
210 if (std::error_code EC = ImageOrErr.getError())
211 handleError(E: errorCodeToError(EC));
212 MemoryBufferRef Image = **ImageOrErr;
213
214 ol_platform_backend_t Backend = OL_PLATFORM_BACKEND_UNKNOWN;
215 ol_init_args_t InitArgs = OL_INIT_ARGS_INIT;
216
217 file_magic Magic = identify_magic(magic: Image.getBuffer());
218 if (Magic >= file_magic::elf && Magic <= file_magic::elf_core) {
219 Expected<object::ELFFile<object::ELF64LE>> ElfOrErr =
220 object::ELFFile<object::ELF64LE>::create(Object: Image.getBuffer());
221 if (!ElfOrErr)
222 handleError(E: ElfOrErr.takeError());
223
224 switch (ElfOrErr->getHeader().e_machine) {
225 case ELF::EM_AMDGPU:
226 Backend = OL_PLATFORM_BACKEND_AMDGPU;
227 break;
228 case ELF::EM_CUDA:
229 Backend = OL_PLATFORM_BACKEND_CUDA;
230 break;
231 default:
232 handleError(E: createStringError(
233 Fmt: "unhandled ELF architecture: %s",
234 Vals: ELF::convertEMachineToArchName(EMachine: ElfOrErr->getHeader().e_machine)
235 .data()));
236 }
237 }
238
239 if (Backend != OL_PLATFORM_BACKEND_UNKNOWN) {
240 InitArgs.NumPlatforms = 1;
241 InitArgs.Platforms = &Backend;
242 }
243
244 SmallVector<const char *> NewArgv = {File.c_str()};
245 llvm::transform(Range&: Args, d_first: std::back_inserter(x&: NewArgv),
246 F: [](const std::string &Arg) { return Arg.c_str(); });
247
248 OFFLOAD_ERR(olInit(&InitArgs));
249 ol_device_handle_t Device = findDevice(Binary: Image);
250 if (!Device)
251 handleError(E: createStringError(Fmt: "No compatible device was found"));
252 ol_device_handle_t Host = getHostDevice();
253 assert(Host && "Host device should always be present");
254
255 ol_context_handle_t Context;
256 OFFLOAD_ERR(olCreateContext(1, &Device, &Context));
257
258 ol_program_handle_t Program;
259 OFFLOAD_ERR(olCreateProgram(Device, Image.getBufferStart(),
260 Image.getBufferSize(), &Program));
261
262 ol_queue_handle_t Queue;
263 OFFLOAD_ERR(olCreateQueue(Context, Device, &Queue));
264
265 int DevArgc = static_cast<int>(NewArgv.size());
266 void *DevArgv = copyArgumentVector(Argc: NewArgv.size(), Argv: NewArgv.begin(), Device);
267 void *DevEnvp = copyEnvironment(Envp: envp, Device);
268
269 void *DevRet;
270 int Zero = 0;
271 OFFLOAD_ERR(olMemAlloc(Device, OL_ALLOC_TYPE_DEVICE, sizeof(int), &DevRet));
272 OFFLOAD_ERR(olMemcpy(Queue, DevRet, Device, &Zero, Host, sizeof(int)));
273
274 uint32_t Dims = (BlocksZ > 1) ? 3 : (BlocksY > 1) ? 2 : 1;
275 ol_kernel_launch_size_args_t StartLaunch{.Dimensions: Dims,
276 .NumGroups: {.x: BlocksX, .y: BlocksY, .z: BlocksZ},
277 .GroupSize: {.x: ThreadsX, .y: ThreadsY, .z: ThreadsZ},
278 /*SharedMemBytes=*/.DynSharedMemory: 0};
279 if (!Kernels.empty()) {
280 // Launch the user-specified kernels in order. These must take no arguments.
281 for (const std::string &Kernel : Kernels)
282 launchKernel(Queue, Device, Program, Name: Kernel.c_str(), LaunchArgs: StartLaunch);
283 } else {
284 // The '_begin' and '_end' kernels perform libc startup and teardown. Global
285 // constructors and destructors are handled automatically by the runtime.
286 ol_kernel_launch_size_args_t BeginLaunch{.Dimensions: 1, .NumGroups: {.x: 1, .y: 1, .z: 1}, .GroupSize: {.x: 1, .y: 1, .z: 1}, .DynSharedMemory: 0};
287 launchKernel(Queue, Device, Program, Name: "_begin", LaunchArgs: BeginLaunch, KernelArgs&: DevArgc,
288 KernelArgs&: DevArgv, KernelArgs&: DevEnvp);
289 OFFLOAD_ERR(olSyncQueue(Queue));
290
291 launchKernel(Queue, Device, Program, Name: "_start", LaunchArgs: StartLaunch, KernelArgs&: DevArgc,
292 KernelArgs&: DevArgv, KernelArgs&: DevEnvp, KernelArgs&: DevRet);
293
294 ol_kernel_launch_size_args_t EndLaunch{.Dimensions: 1, .NumGroups: {.x: 1, .y: 1, .z: 1}, .GroupSize: {.x: 1, .y: 1, .z: 1}, .DynSharedMemory: 0};
295 launchKernel(Queue, Device, Program, Name: "_end", LaunchArgs: EndLaunch);
296 }
297
298 int Ret;
299 OFFLOAD_ERR(olMemcpy(Queue, &Ret, Host, DevRet, Device, sizeof(int)));
300 OFFLOAD_ERR(olSyncQueue(Queue));
301
302 OFFLOAD_ERR(olMemFree(DevRet));
303 OFFLOAD_ERR(olMemFree(DevArgv));
304 OFFLOAD_ERR(olMemFree(DevEnvp));
305 OFFLOAD_ERR(olDestroyQueue(Queue));
306 OFFLOAD_ERR(olDestroyContext(Context));
307 OFFLOAD_ERR(olDestroyProgram(Program));
308 OFFLOAD_ERR(olShutDown());
309
310 return Ret;
311}
312