368 lines · cpp
1//===-- Loader Implementation for NVPTX 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-exception6//7//===----------------------------------------------------------------------===//8//9// This file impelements a simple loader to run images supporting the NVPTX10// architecture. The file launches the '_start' kernel which should be provided11// by the device application start code and call ultimately call the 'main'12// function.13//14//===----------------------------------------------------------------------===//15 16#include "llvm-gpu-loader.h"17#include "server.h"18 19#include "cuda.h"20 21#include "llvm/Object/ELF.h"22#include "llvm/Object/ELFObjectFile.h"23 24#include <atomic>25#include <cstddef>26#include <cstdio>27#include <cstdlib>28#include <cstring>29#include <thread>30#include <vector>31 32using namespace llvm;33using namespace object;34 35static void handle_error_impl(const char *file, int32_t line, CUresult err) {36 if (err == CUDA_SUCCESS)37 return;38 39 const char *err_str = nullptr;40 CUresult result = cuGetErrorString(err, &err_str);41 if (result != CUDA_SUCCESS)42 fprintf(stderr, "%s:%d:0: Unknown Error\n", file, line);43 else44 fprintf(stderr, "%s:%d:0: Error: %s\n", file, line, err_str);45 exit(1);46}47 48// Gets the names of all the globals that contain functions to initialize or49// deinitialize. We need to do this manually because the NVPTX toolchain does50// not contain the necessary binary manipulation tools.51template <typename Alloc>52Expected<void *> get_ctor_dtor_array(const void *image, const size_t size,53 Alloc allocator, CUmodule binary) {54 auto mem_buffer = MemoryBuffer::getMemBuffer(55 StringRef(reinterpret_cast<const char *>(image), size), "image",56 /*RequiresNullTerminator=*/false);57 Expected<ELF64LEObjectFile> elf_or_err =58 ELF64LEObjectFile::create(*mem_buffer);59 if (!elf_or_err)60 handle_error(toString(elf_or_err.takeError()).c_str());61 62 std::vector<std::pair<const char *, uint16_t>> ctors;63 std::vector<std::pair<const char *, uint16_t>> dtors;64 // CUDA has no way to iterate over all the symbols so we need to inspect the65 // ELF directly using the LLVM libraries.66 for (const auto &symbol : elf_or_err->symbols()) {67 auto name_or_err = symbol.getName();68 if (!name_or_err)69 handle_error(toString(name_or_err.takeError()).c_str());70 71 // Search for all symbols that contain a constructor or destructor.72 if (!name_or_err->starts_with("__init_array_object_") &&73 !name_or_err->starts_with("__fini_array_object_"))74 continue;75 76 uint16_t priority;77 if (name_or_err->rsplit('_').second.getAsInteger(10, priority))78 handle_error("Invalid priority for constructor or destructor");79 80 if (name_or_err->starts_with("__init"))81 ctors.emplace_back(std::make_pair(name_or_err->data(), priority));82 else83 dtors.emplace_back(std::make_pair(name_or_err->data(), priority));84 }85 // Lower priority constructors are run before higher ones. The reverse is true86 // for destructors.87 llvm::sort(ctors, llvm::less_second());88 llvm::sort(dtors, llvm::less_second());89 90 // Allocate host pinned memory to make these arrays visible to the GPU.91 CUdeviceptr *dev_memory = reinterpret_cast<CUdeviceptr *>(allocator(92 ctors.size() * sizeof(CUdeviceptr) + dtors.size() * sizeof(CUdeviceptr)));93 uint64_t global_size = 0;94 95 // Get the address of the global and then store the address of the constructor96 // function to call in the constructor array.97 CUdeviceptr *dev_ctors_start = dev_memory;98 CUdeviceptr *dev_ctors_end = dev_ctors_start + ctors.size();99 for (uint64_t i = 0; i < ctors.size(); ++i) {100 CUdeviceptr dev_ptr;101 if (CUresult err =102 cuModuleGetGlobal(&dev_ptr, &global_size, binary, ctors[i].first))103 handle_error(err);104 if (CUresult err =105 cuMemcpyDtoH(&dev_ctors_start[i], dev_ptr, sizeof(uintptr_t)))106 handle_error(err);107 }108 109 // Get the address of the global and then store the address of the destructor110 // function to call in the destructor array.111 CUdeviceptr *dev_dtors_start = dev_ctors_end;112 CUdeviceptr *dev_dtors_end = dev_dtors_start + dtors.size();113 for (uint64_t i = 0; i < dtors.size(); ++i) {114 CUdeviceptr dev_ptr;115 if (CUresult err =116 cuModuleGetGlobal(&dev_ptr, &global_size, binary, dtors[i].first))117 handle_error(err);118 if (CUresult err =119 cuMemcpyDtoH(&dev_dtors_start[i], dev_ptr, sizeof(uintptr_t)))120 handle_error(err);121 }122 123 // Obtain the address of the pointers the startup implementation uses to124 // iterate the constructors and destructors.125 CUdeviceptr init_start;126 if (CUresult err = cuModuleGetGlobal(&init_start, &global_size, binary,127 "__init_array_start"))128 handle_error(err);129 CUdeviceptr init_end;130 if (CUresult err = cuModuleGetGlobal(&init_end, &global_size, binary,131 "__init_array_end"))132 handle_error(err);133 CUdeviceptr fini_start;134 if (CUresult err = cuModuleGetGlobal(&fini_start, &global_size, binary,135 "__fini_array_start"))136 handle_error(err);137 CUdeviceptr fini_end;138 if (CUresult err = cuModuleGetGlobal(&fini_end, &global_size, binary,139 "__fini_array_end"))140 handle_error(err);141 142 // Copy the pointers to the newly written array to the symbols so the startup143 // implementation can iterate them.144 if (CUresult err =145 cuMemcpyHtoD(init_start, &dev_ctors_start, sizeof(uintptr_t)))146 handle_error(err);147 if (CUresult err = cuMemcpyHtoD(init_end, &dev_ctors_end, sizeof(uintptr_t)))148 handle_error(err);149 if (CUresult err =150 cuMemcpyHtoD(fini_start, &dev_dtors_start, sizeof(uintptr_t)))151 handle_error(err);152 if (CUresult err = cuMemcpyHtoD(fini_end, &dev_dtors_end, sizeof(uintptr_t)))153 handle_error(err);154 155 return dev_memory;156}157 158void print_kernel_resources(CUmodule binary, const char *kernel_name) {159 CUfunction function;160 if (CUresult err = cuModuleGetFunction(&function, binary, kernel_name))161 handle_error(err);162 int num_regs;163 if (CUresult err =164 cuFuncGetAttribute(&num_regs, CU_FUNC_ATTRIBUTE_NUM_REGS, function))165 handle_error(err);166 printf("Executing kernel %s:\n", kernel_name);167 printf("%6s registers: %d\n", kernel_name, num_regs);168}169 170template <typename args_t>171CUresult launch_kernel(CUmodule binary, CUstream stream, rpc::Server &server,172 const LaunchParameters ¶ms, const char *kernel_name,173 args_t kernel_args, bool print_resource_usage) {174 // look up the '_start' kernel in the loaded module.175 CUfunction function;176 if (CUresult err = cuModuleGetFunction(&function, binary, kernel_name))177 handle_error(err);178 179 // Set up the arguments to the '_start' kernel on the GPU.180 uint64_t args_size = std::is_empty_v<args_t> ? 0 : sizeof(args_t);181 void *args_config[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, &kernel_args,182 CU_LAUNCH_PARAM_BUFFER_SIZE, &args_size,183 CU_LAUNCH_PARAM_END};184 if (print_resource_usage)185 print_kernel_resources(binary, kernel_name);186 187 // Initialize a non-blocking CUDA stream to allocate memory if needed.188 // This needs to be done on a separate stream or else it will deadlock189 // with the executing kernel.190 CUstream memory_stream;191 if (CUresult err = cuStreamCreate(&memory_stream, CU_STREAM_NON_BLOCKING))192 handle_error(err);193 194 std::atomic<bool> finished = false;195 std::thread server_thread(196 [](std::atomic<bool> *finished, rpc::Server *server,197 CUstream memory_stream) {198 auto malloc_handler = [&](size_t size) -> void * {199 CUdeviceptr dev_ptr;200 if (CUresult err = cuMemAllocAsync(&dev_ptr, size, memory_stream))201 dev_ptr = 0UL;202 203 // Wait until the memory allocation is complete.204 while (cuStreamQuery(memory_stream) == CUDA_ERROR_NOT_READY)205 ;206 return reinterpret_cast<void *>(dev_ptr);207 };208 209 auto free_handler = [&](void *ptr) -> void {210 if (CUresult err = cuMemFreeAsync(reinterpret_cast<CUdeviceptr>(ptr),211 memory_stream))212 handle_error(err);213 };214 215 uint32_t index = 0;216 while (!*finished) {217 index =218 handle_server<32>(*server, index, malloc_handler, free_handler);219 }220 },221 &finished, &server, memory_stream);222 223 // Call the kernel with the given arguments.224 if (CUresult err = cuLaunchKernel(225 function, params.num_blocks_x, params.num_blocks_y,226 params.num_blocks_z, params.num_threads_x, params.num_threads_y,227 params.num_threads_z, 0, stream, nullptr, args_config))228 handle_error(err);229 230 if (CUresult err = cuStreamSynchronize(stream))231 handle_error(err);232 233 finished = true;234 if (server_thread.joinable())235 server_thread.join();236 237 return CUDA_SUCCESS;238}239 240int load_nvptx(int argc, const char **argv, const char **envp, void *image,241 size_t size, const LaunchParameters ¶ms,242 bool print_resource_usage) {243 if (CUresult err = cuInit(0))244 handle_error(err);245 // Obtain the first device found on the system.246 uint32_t device_id = 0;247 CUdevice device;248 if (CUresult err = cuDeviceGet(&device, device_id))249 handle_error(err);250 251 // Initialize the CUDA context and claim it for this execution.252 CUcontext context;253 if (CUresult err = cuDevicePrimaryCtxRetain(&context, device))254 handle_error(err);255 if (CUresult err = cuCtxSetCurrent(context))256 handle_error(err);257 258 // Increase the stack size per thread.259 // TODO: We should allow this to be passed in so only the tests that require a260 // larger stack can specify it to save on memory usage.261 if (CUresult err = cuCtxSetLimit(CU_LIMIT_STACK_SIZE, 3 * 1024))262 handle_error(err);263 264 // Initialize a non-blocking CUDA stream to execute the kernel.265 CUstream stream;266 if (CUresult err = cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING))267 handle_error(err);268 269 // Load the image into a CUDA module.270 CUmodule binary;271 if (CUresult err = cuModuleLoadDataEx(&binary, image, 0, nullptr, nullptr))272 handle_error(err);273 274 // Allocate pinned memory on the host to hold the pointer array for the275 // copied argv and allow the GPU device to access it.276 auto allocator = [&](uint64_t size) -> void * {277 void *dev_ptr;278 if (CUresult err = cuMemAllocHost(&dev_ptr, size))279 handle_error(err);280 return dev_ptr;281 };282 283 auto memory_or_err = get_ctor_dtor_array(image, size, allocator, binary);284 if (!memory_or_err)285 handle_error(toString(memory_or_err.takeError()).c_str());286 287 void *dev_argv = copy_argument_vector(argc, argv, allocator);288 if (!dev_argv)289 handle_error("Failed to allocate device argv");290 291 // Allocate pinned memory on the host to hold the pointer array for the292 // copied environment array and allow the GPU device to access it.293 void *dev_envp = copy_environment(envp, allocator);294 if (!dev_envp)295 handle_error("Failed to allocate device environment");296 297 // Allocate space for the return pointer and initialize it to zero.298 CUdeviceptr dev_ret;299 if (CUresult err = cuMemAlloc(&dev_ret, sizeof(int)))300 handle_error(err);301 if (CUresult err = cuMemsetD32(dev_ret, 0, 1))302 handle_error(err);303 304 uint32_t warp_size = 32;305 void *rpc_buffer = nullptr;306 if (CUresult err = cuMemAllocHost(307 &rpc_buffer,308 rpc::Server::allocation_size(warp_size, rpc::MAX_PORT_COUNT)))309 handle_error(err);310 rpc::Server server(rpc::MAX_PORT_COUNT, rpc_buffer);311 rpc::Client client(rpc::MAX_PORT_COUNT, rpc_buffer);312 313 // Initialize the RPC client on the device by copying the local data to the314 // device's internal pointer.315 CUdeviceptr rpc_client_dev = 0;316 uint64_t client_ptr_size = sizeof(void *);317 if (CUresult err = cuModuleGetGlobal(&rpc_client_dev, &client_ptr_size,318 binary, "__llvm_rpc_client"))319 handle_error(err);320 321 if (CUresult err = cuMemcpyHtoD(rpc_client_dev, &client, sizeof(rpc::Client)))322 handle_error(err);323 324 LaunchParameters single_threaded_params = {1, 1, 1, 1, 1, 1};325 begin_args_t init_args = {argc, dev_argv, dev_envp};326 if (CUresult err =327 launch_kernel(binary, stream, server, single_threaded_params,328 "_begin", init_args, print_resource_usage))329 handle_error(err);330 331 start_args_t args = {argc, dev_argv, dev_envp,332 reinterpret_cast<void *>(dev_ret)};333 if (CUresult err = launch_kernel(binary, stream, server, params, "_start",334 args, print_resource_usage))335 handle_error(err);336 337 // Copy the return value back from the kernel and wait.338 int host_ret = 0;339 if (CUresult err = cuMemcpyDtoH(&host_ret, dev_ret, sizeof(int)))340 handle_error(err);341 342 if (CUresult err = cuStreamSynchronize(stream))343 handle_error(err);344 345 end_args_t fini_args = {};346 if (CUresult err =347 launch_kernel(binary, stream, server, single_threaded_params, "_end",348 fini_args, print_resource_usage))349 handle_error(err);350 351 // Free the memory allocated for the device.352 if (CUresult err = cuMemFreeHost(*memory_or_err))353 handle_error(err);354 if (CUresult err = cuMemFree(dev_ret))355 handle_error(err);356 if (CUresult err = cuMemFreeHost(dev_argv))357 handle_error(err);358 if (CUresult err = cuMemFreeHost(rpc_buffer))359 handle_error(err);360 361 // Destroy the context and the loaded binary.362 if (CUresult err = cuModuleUnload(binary))363 handle_error(err);364 if (CUresult err = cuDevicePrimaryCtxRelease(device))365 handle_error(err);366 return host_ret;367}368