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