233 lines · cpp
1//===- RPC.h - Interface for remote procedure calls from the GPU ----------===//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#include "RPC.h"10 11#include "Shared/Debug.h"12#include "Shared/RPCOpcodes.h"13 14#include "PluginInterface.h"15 16#include "shared/rpc.h"17#include "shared/rpc_opcodes.h"18#include "shared/rpc_server.h"19 20using namespace llvm;21using namespace omp;22using namespace target;23 24template <uint32_t NumLanes>25rpc::Status handleOffloadOpcodes(plugin::GenericDeviceTy &Device,26 rpc::Server::Port &Port) {27 28 switch (Port.get_opcode()) {29 case LIBC_MALLOC: {30 Port.recv_and_send([&](rpc::Buffer *Buffer, uint32_t) {31 auto PtrOrErr =32 Device.allocate(Buffer->data[0], nullptr, TARGET_ALLOC_DEVICE);33 void *Ptr = nullptr;34 if (!PtrOrErr)35 llvm::consumeError(PtrOrErr.takeError());36 else37 Ptr = *PtrOrErr;38 Buffer->data[0] = reinterpret_cast<uintptr_t>(Ptr);39 });40 break;41 }42 case LIBC_FREE: {43 Port.recv([&](rpc::Buffer *Buffer, uint32_t) {44 if (auto Err = Device.free(reinterpret_cast<void *>(Buffer->data[0]),45 TARGET_ALLOC_DEVICE))46 llvm::consumeError(std::move(Err));47 });48 break;49 }50 case OFFLOAD_HOST_CALL: {51 uint64_t Sizes[NumLanes] = {0};52 unsigned long long Results[NumLanes] = {0};53 void *Args[NumLanes] = {nullptr};54 Port.recv_n(Args, Sizes, [&](uint64_t Size) { return new char[Size]; });55 Port.recv([&](rpc::Buffer *buffer, uint32_t ID) {56 using FuncPtrTy = unsigned long long (*)(void *);57 auto Func = reinterpret_cast<FuncPtrTy>(buffer->data[0]);58 Results[ID] = Func(Args[ID]);59 });60 Port.send([&](rpc::Buffer *Buffer, uint32_t ID) {61 Buffer->data[0] = static_cast<uint64_t>(Results[ID]);62 delete[] reinterpret_cast<char *>(Args[ID]);63 });64 break;65 }66 default:67 return rpc::RPC_UNHANDLED_OPCODE;68 break;69 }70 return rpc::RPC_SUCCESS;71}72 73static rpc::Status handleOffloadOpcodes(plugin::GenericDeviceTy &Device,74 rpc::Server::Port &Port,75 uint32_t NumLanes) {76 if (NumLanes == 1)77 return handleOffloadOpcodes<1>(Device, Port);78 else if (NumLanes == 32)79 return handleOffloadOpcodes<32>(Device, Port);80 else if (NumLanes == 64)81 return handleOffloadOpcodes<64>(Device, Port);82 else83 return rpc::RPC_ERROR;84}85 86static rpc::Status runServer(plugin::GenericDeviceTy &Device, void *Buffer,87 bool &ClientInUse) {88 uint64_t NumPorts =89 std::min(Device.requestedRPCPortCount(), rpc::MAX_PORT_COUNT);90 rpc::Server Server(NumPorts, Buffer);91 92 auto Port = Server.try_open(Device.getWarpSize());93 if (!Port)94 return rpc::RPC_SUCCESS;95 96 ClientInUse = true;97 rpc::Status Status =98 handleOffloadOpcodes(Device, *Port, Device.getWarpSize());99 100 // Let the `libc` library handle any other unhandled opcodes.101 if (Status == rpc::RPC_UNHANDLED_OPCODE)102 Status = LIBC_NAMESPACE::shared::handle_libc_opcodes(*Port,103 Device.getWarpSize());104 Port->close();105 106 return Status;107}108 109void RPCServerTy::ServerThread::startThread() {110 if (!Running.fetch_or(true, std::memory_order_acquire))111 Worker = std::thread([this]() { run(); });112}113 114void RPCServerTy::ServerThread::shutDown() {115 if (!Running.fetch_and(false, std::memory_order_release))116 return;117 {118 std::lock_guard<decltype(Mutex)> Lock(Mutex);119 CV.notify_all();120 }121 if (Worker.joinable())122 Worker.join();123}124 125void RPCServerTy::ServerThread::run() {126 static constexpr auto IdleTime = std::chrono::microseconds(25);127 static constexpr auto IdleSleep = std::chrono::microseconds(250);128 std::unique_lock<decltype(Mutex)> Lock(Mutex);129 130 auto LastUse = std::chrono::steady_clock::now();131 for (;;) {132 CV.wait(Lock, [&]() {133 return NumUsers.load(std::memory_order_acquire) > 0 ||134 !Running.load(std::memory_order_acquire);135 });136 137 if (!Running.load(std::memory_order_acquire))138 return;139 140 Lock.unlock();141 bool ClientInUse = false;142 while (NumUsers.load(std::memory_order_relaxed) > 0 &&143 Running.load(std::memory_order_relaxed)) {144 145 // Suspend this thread briefly if there is no current work.146 auto Now = std::chrono::steady_clock::now();147 if (!ClientInUse && Now - LastUse >= IdleTime)148 std::this_thread::sleep_for(IdleSleep);149 else if (ClientInUse)150 LastUse = Now;151 152 ClientInUse = false;153 std::lock_guard<decltype(Mutex)> Lock(BufferMutex);154 for (const auto &[Buffer, Device] : llvm::zip_equal(Buffers, Devices)) {155 if (!Buffer || !Device)156 continue;157 158 // If running the server failed, print a message but keep running.159 if (runServer(*Device, Buffer, ClientInUse) != rpc::RPC_SUCCESS)160 FAILURE_MESSAGE("Unhandled or invalid RPC opcode!");161 }162 }163 Lock.lock();164 }165}166 167RPCServerTy::RPCServerTy(plugin::GenericPluginTy &Plugin)168 : Buffers(std::make_unique<void *[]>(Plugin.getNumDevices())),169 Devices(std::make_unique<plugin::GenericDeviceTy *[]>(170 Plugin.getNumDevices())),171 Thread(new ServerThread(Buffers.get(), Devices.get(),172 Plugin.getNumDevices(), BufferMutex)) {}173 174llvm::Error RPCServerTy::startThread() {175 Thread->startThread();176 return Error::success();177}178 179llvm::Error RPCServerTy::shutDown() {180 Thread->shutDown();181 return Error::success();182}183 184llvm::Expected<bool>185RPCServerTy::isDeviceUsingRPC(plugin::GenericDeviceTy &Device,186 plugin::GenericGlobalHandlerTy &Handler,187 plugin::DeviceImageTy &Image) {188 return Handler.isSymbolInImage(Device, Image, "__llvm_rpc_client");189}190 191Error RPCServerTy::initDevice(plugin::GenericDeviceTy &Device,192 plugin::GenericGlobalHandlerTy &Handler,193 plugin::DeviceImageTy &Image) {194 uint64_t NumPorts =195 std::min(Device.requestedRPCPortCount(), rpc::MAX_PORT_COUNT);196 auto RPCBufferOrErr = Device.allocate(197 rpc::Server::allocation_size(Device.getWarpSize(), NumPorts), nullptr,198 TARGET_ALLOC_HOST);199 if (!RPCBufferOrErr)200 return RPCBufferOrErr.takeError();201 202 void *RPCBuffer = *RPCBufferOrErr;203 if (!RPCBuffer)204 return plugin::Plugin::error(205 error::ErrorCode::UNKNOWN,206 "failed to initialize RPC server for device %d", Device.getDeviceId());207 208 // Get the address of the RPC client from the device.209 plugin::GlobalTy ClientGlobal("__llvm_rpc_client", sizeof(rpc::Client));210 if (auto Err =211 Handler.getGlobalMetadataFromDevice(Device, Image, ClientGlobal))212 return Err;213 214 rpc::Client client(NumPorts, RPCBuffer);215 if (auto Err = Device.dataSubmit(ClientGlobal.getPtr(), &client,216 sizeof(rpc::Client), nullptr))217 return Err;218 std::lock_guard<decltype(BufferMutex)> Lock(BufferMutex);219 Buffers[Device.getDeviceId()] = RPCBuffer;220 Devices[Device.getDeviceId()] = &Device;221 222 return Error::success();223}224 225Error RPCServerTy::deinitDevice(plugin::GenericDeviceTy &Device) {226 std::lock_guard<decltype(BufferMutex)> Lock(BufferMutex);227 if (auto Err = Device.free(Buffers[Device.getDeviceId()], TARGET_ALLOC_HOST))228 return Err;229 Buffers[Device.getDeviceId()] = nullptr;230 Devices[Device.getDeviceId()] = nullptr;231 return Error::success();232}233