openmp/libomptarget/plugins/remote/server/Server.cpp

   1 //===----------------- Server.cpp - Server Implementation -----------------===//
   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 // Offloading gRPC server for remote host.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include <cmath>
  14 #include <future>
  15
  16 #include "Server.h"
  17 #include "omptarget.h"
  18 #include "openmp.grpc.pb.h"
  19 #include "openmp.pb.h"
  20
  21 using grpc::WriteOptions;
  22
  23 extern std::promise<void> ShutdownPromise;
  24
  25 Status RemoteOffloadImpl::Shutdown(ServerContext *Context, const Null *Request,
  26                                    I32 *Reply) {
  27   SERVER_DBG("Shutting down the server")
  28
  29   Reply->set_number(0);
  30   ShutdownPromise.set_value();
  31   return Status::OK;
  32 }
  33
  34 Status
  35 RemoteOffloadImpl::RegisterLib(ServerContext *Context,
  36                                const TargetBinaryDescription *Description,
  37                                I32 *Reply) {
  38   auto Desc = std::make_unique<__tgt_bin_desc>();
  39
  40   unloadTargetBinaryDescription(Description, Desc.get(),
  41                                 HostToRemoteDeviceImage);
  42   PM->RTLs.RegisterLib(Desc.get());
  43
  44   if (Descriptions.find((void *)Description->bin_ptr()) != Descriptions.end())
  45     freeTargetBinaryDescription(
  46         Descriptions[(void *)Description->bin_ptr()].get());
  47   else
  48     Descriptions[(void *)Description->bin_ptr()] = std::move(Desc);
  49
  50   SERVER_DBG("Registered library")
  51   Reply->set_number(0);
  52   return Status::OK;
  53 }
  54
  55 Status RemoteOffloadImpl::UnregisterLib(ServerContext *Context,
  56                                         const Pointer *Request, I32 *Reply) {
  57   if (Descriptions.find((void *)Request->number()) == Descriptions.end()) {
  58     Reply->set_number(1);
  59     return Status::OK;
  60   }
  61
  62   PM->RTLs.UnregisterLib(Descriptions[(void *)Request->number()].get());
  63   freeTargetBinaryDescription(Descriptions[(void *)Request->number()].get());
  64   Descriptions.erase((void *)Request->number());
  65
  66   SERVER_DBG("Unregistered library")
  67   Reply->set_number(0);
  68   return Status::OK;
  69 }
  70
  71 Status RemoteOffloadImpl::IsValidBinary(ServerContext *Context,
  72                                         const TargetDeviceImagePtr *DeviceImage,
  73                                         I32 *IsValid) {
  74   __tgt_device_image *Image =
  75       HostToRemoteDeviceImage[(void *)DeviceImage->image_ptr()];
  76
  77   IsValid->set_number(0);
  78
  79   for (auto &RTL : PM->RTLs.AllRTLs)
  80     if (auto Ret = RTL.is_valid_binary(Image)) {
  81       IsValid->set_number(Ret);
  82       break;
  83     }
  84
  85   SERVER_DBG("Checked if binary (%p) is valid",
  86              (void *)(DeviceImage->image_ptr()))
  87   return Status::OK;
  88 }
  89
  90 Status RemoteOffloadImpl::GetNumberOfDevices(ServerContext *Context,
  91                                              const Null *Null,
  92                                              I32 *NumberOfDevices) {
  93   std::call_once(PM->RTLs.initFlag, &RTLsTy::LoadRTLs, &PM->RTLs);
  94
  95   int32_t Devices = 0;
  96   PM->RTLsMtx.lock();
  97   for (auto &RTL : PM->RTLs.AllRTLs)
  98     Devices += RTL.NumberOfDevices;
  99   PM->RTLsMtx.unlock();
 100
 101   NumberOfDevices->set_number(Devices);
 102
 103   SERVER_DBG("Got number of devices")
 104   return Status::OK;
 105 }
 106
 107 Status RemoteOffloadImpl::InitDevice(ServerContext *Context,
 108                                      const I32 *DeviceNum, I32 *Reply) {
 109   Reply->set_number(PM->Devices[DeviceNum->number()].RTL->init_device(
 110       mapHostRTLDeviceId(DeviceNum->number())));
 111
 112   SERVER_DBG("Initialized device %d", DeviceNum->number())
 113   return Status::OK;
 114 }
 115
 116 Status RemoteOffloadImpl::InitRequires(ServerContext *Context,
 117                                        const I64 *RequiresFlag, I32 *Reply) {
 118   for (auto &Device : PM->Devices)
 119     if (Device.RTL->init_requires)
 120       Device.RTL->init_requires(RequiresFlag->number());
 121   Reply->set_number(RequiresFlag->number());
 122
 123   SERVER_DBG("Initialized requires for devices")
 124   return Status::OK;
 125 }
 126
 127 Status RemoteOffloadImpl::LoadBinary(ServerContext *Context,
 128                                      const Binary *Binary, TargetTable *Reply) {
 129   __tgt_device_image *Image =
 130       HostToRemoteDeviceImage[(void *)Binary->image_ptr()];
 131
 132   Table = PM->Devices[Binary->device_id()].RTL->load_binary(
 133       mapHostRTLDeviceId(Binary->device_id()), Image);
 134   if (Table)
 135     loadTargetTable(Table, *Reply, Image);
 136
 137   SERVER_DBG("Loaded binary (%p) to device %d", (void *)Binary->image_ptr(),
 138              Binary->device_id())
 139   return Status::OK;
 140 }
 141
 142 Status RemoteOffloadImpl::IsDataExchangeable(ServerContext *Context,
 143                                              const DevicePair *Request,
 144                                              I32 *Reply) {
 145   Reply->set_number(-1);
 146   if (PM->Devices[mapHostRTLDeviceId(Request->src_dev_id())]
 147           .RTL->is_data_exchangable)
 148     Reply->set_number(PM->Devices[mapHostRTLDeviceId(Request->src_dev_id())]
 149                           .RTL->is_data_exchangable(Request->src_dev_id(),
 150                                                     Request->dst_dev_id()));
 151
 152   SERVER_DBG("Checked if data exchangeable between device %d and device %d",
 153              Request->src_dev_id(), Request->dst_dev_id())
 154   return Status::OK;
 155 }
 156
 157 Status RemoteOffloadImpl::DataAlloc(ServerContext *Context,
 158                                     const AllocData *Request, Pointer *Reply) {
 159   uint64_t TgtPtr = (uint64_t)PM->Devices[Request->device_id()].RTL->data_alloc(
 160       mapHostRTLDeviceId(Request->device_id()), Request->size(),
 161       (void *)Request->hst_ptr(), TARGET_ALLOC_DEFAULT);
 162   Reply->set_number(TgtPtr);
 163
 164   SERVER_DBG("Allocated at " DPxMOD "", DPxPTR((void *)TgtPtr))
 165
 166   return Status::OK;
 167 }
 168
 169 Status RemoteOffloadImpl::DataSubmit(ServerContext *Context,
 170                                      ServerReader<SubmitData> *Reader,
 171                                      I32 *Reply) {
 172   SubmitData Request;
 173   uint8_t *HostCopy = nullptr;
 174   while (Reader->Read(&Request)) {
 175     if (Request.start() == 0 && Request.size() == Request.data().size()) {
 176       Reader->SendInitialMetadata();
 177
 178       Reply->set_number(PM->Devices[Request.device_id()].RTL->data_submit(
 179           mapHostRTLDeviceId(Request.device_id()), (void *)Request.tgt_ptr(),
 180           (void *)Request.data().data(), Request.data().size()));
 181
 182       SERVER_DBG("Submitted %lu bytes async to (%p) on device %d",
 183                  Request.data().size(), (void *)Request.tgt_ptr(),
 184                  Request.device_id())
 185
 186       return Status::OK;
 187     }
 188     if (!HostCopy) {
 189       HostCopy = new uint8_t[Request.size()];
 190       Reader->SendInitialMetadata();
 191     }
 192
 193     memcpy((void *)((char *)HostCopy + Request.start()), Request.data().data(),
 194            Request.data().size());
 195   }
 196
 197   Reply->set_number(PM->Devices[Request.device_id()].RTL->data_submit(
 198       mapHostRTLDeviceId(Request.device_id()), (void *)Request.tgt_ptr(),
 199       HostCopy, Request.size()));
 200
 201   delete[] HostCopy;
 202
 203   SERVER_DBG("Submitted %lu bytes to (%p) on device %d", Request.data().size(),
 204              (void *)Request.tgt_ptr(), Request.device_id())
 205
 206   return Status::OK;
 207 }
 208
 209 Status RemoteOffloadImpl::DataRetrieve(ServerContext *Context,
 210                                        const RetrieveData *Request,
 211                                        ServerWriter<Data> *Writer) {
 212   auto HstPtr = std::make_unique<char[]>(Request->size());
 213
 214   auto Ret = PM->Devices[Request->device_id()].RTL->data_retrieve(
 215       mapHostRTLDeviceId(Request->device_id()), HstPtr.get(),
 216       (void *)Request->tgt_ptr(), Request->size());
 217
 218   if (Arena->SpaceAllocated() >= MaxSize)
 219     Arena->Reset();
 220
 221   if (Request->size() > BlockSize) {
 222     uint64_t Start = 0, End = BlockSize;
 223     for (auto I = 0; I < ceil((float)Request->size() / BlockSize); I++) {
 224       auto *Reply = protobuf::Arena::CreateMessage<Data>(Arena.get());
 225
 226       Reply->set_start(Start);
 227       Reply->set_size(Request->size());
 228       Reply->set_data((char *)HstPtr.get() + Start, End - Start);
 229       Reply->set_ret(Ret);
 230
 231       if (!Writer->Write(*Reply)) {
 232         CLIENT_DBG("Broken stream when submitting data")
 233       }
 234
 235       SERVER_DBG("Retrieved %lu-%lu/%lu bytes from (%p) on device %d", Start,
 236                  End, Request->size(), (void *)Request->tgt_ptr(),
 237                  mapHostRTLDeviceId(Request->device_id()))
 238
 239       Start += BlockSize;
 240       End += BlockSize;
 241       if (End >= Request->size())
 242         End = Request->size();
 243     }
 244   } else {
 245     auto *Reply = protobuf::Arena::CreateMessage<Data>(Arena.get());
 246
 247     Reply->set_start(0);
 248     Reply->set_size(Request->size());
 249     Reply->set_data((char *)HstPtr.get(), Request->size());
 250     Reply->set_ret(Ret);
 251
 252     SERVER_DBG("Retrieved %lu bytes from (%p) on device %d", Request->size(),
 253                (void *)Request->tgt_ptr(),
 254                mapHostRTLDeviceId(Request->device_id()))
 255
 256     Writer->WriteLast(*Reply, WriteOptions());
 257   }
 258
 259   return Status::OK;
 260 }
 261
 262 Status RemoteOffloadImpl::DataExchange(ServerContext *Context,
 263                                        const ExchangeData *Request,
 264                                        I32 *Reply) {
 265   if (PM->Devices[Request->src_dev_id()].RTL->data_exchange) {
 266     int32_t Ret = PM->Devices[Request->src_dev_id()].RTL->data_exchange(
 267         mapHostRTLDeviceId(Request->src_dev_id()), (void *)Request->src_ptr(),
 268         mapHostRTLDeviceId(Request->dst_dev_id()), (void *)Request->dst_ptr(),
 269         Request->size());
 270     Reply->set_number(Ret);
 271   } else
 272     Reply->set_number(-1);
 273
 274   SERVER_DBG(
 275       "Exchanged data asynchronously from device %d (%p) to device %d (%p) of "
 276       "size %lu",
 277       mapHostRTLDeviceId(Request->src_dev_id()), (void *)Request->src_ptr(),
 278       mapHostRTLDeviceId(Request->dst_dev_id()), (void *)Request->dst_ptr(),
 279       Request->size())
 280   return Status::OK;
 281 }
 282
 283 Status RemoteOffloadImpl::DataDelete(ServerContext *Context,
 284                                      const DeleteData *Request, I32 *Reply) {
 285   auto Ret = PM->Devices[Request->device_id()].RTL->data_delete(
 286       mapHostRTLDeviceId(Request->device_id()), (void *)Request->tgt_ptr());
 287   Reply->set_number(Ret);
 288
 289   SERVER_DBG("Deleted data from (%p) on device %d", (void *)Request->tgt_ptr(),
 290              mapHostRTLDeviceId(Request->device_id()))
 291   return Status::OK;
 292 }
 293
 294 Status RemoteOffloadImpl::RunTargetRegion(ServerContext *Context,
 295                                           const TargetRegion *Request,
 296                                           I32 *Reply) {
 297   std::vector<uint8_t> TgtArgs(Request->arg_num());
 298   for (auto I = 0; I < Request->arg_num(); I++)
 299     TgtArgs[I] = (uint64_t)Request->tgt_args()[I];
 300
 301   std::vector<ptrdiff_t> TgtOffsets(Request->arg_num());
 302   const auto *TgtOffsetItr = Request->tgt_offsets().begin();
 303   for (auto I = 0; I < Request->arg_num(); I++, TgtOffsetItr++)
 304     TgtOffsets[I] = (ptrdiff_t)*TgtOffsetItr;
 305
 306   void *TgtEntryPtr = ((__tgt_offload_entry *)Request->tgt_entry_ptr())->addr;
 307
 308   int32_t Ret = PM->Devices[Request->device_id()].RTL->run_region(
 309       mapHostRTLDeviceId(Request->device_id()), TgtEntryPtr,
 310       (void **)TgtArgs.data(), TgtOffsets.data(), Request->arg_num());
 311
 312   Reply->set_number(Ret);
 313
 314   SERVER_DBG("Ran TargetRegion on device %d with %d args",
 315              mapHostRTLDeviceId(Request->device_id()), Request->arg_num())
 316   return Status::OK;
 317 }
 318
 319 Status RemoteOffloadImpl::RunTargetTeamRegion(ServerContext *Context,
 320                                               const TargetTeamRegion *Request,
 321                                               I32 *Reply) {
 322   std::vector<uint64_t> TgtArgs(Request->arg_num());
 323   for (auto I = 0; I < Request->arg_num(); I++)
 324     TgtArgs[I] = (uint64_t)Request->tgt_args()[I];
 325
 326   std::vector<ptrdiff_t> TgtOffsets(Request->arg_num());
 327   const auto *TgtOffsetItr = Request->tgt_offsets().begin();
 328   for (auto I = 0; I < Request->arg_num(); I++, TgtOffsetItr++)
 329     TgtOffsets[I] = (ptrdiff_t)*TgtOffsetItr;
 330
 331   void *TgtEntryPtr = ((__tgt_offload_entry *)Request->tgt_entry_ptr())->addr;
 332
 333   int32_t Ret = PM->Devices[Request->device_id()].RTL->run_team_region(
 334       mapHostRTLDeviceId(Request->device_id()), TgtEntryPtr,
 335       (void **)TgtArgs.data(), TgtOffsets.data(), Request->arg_num(),
 336       Request->team_num(), Request->thread_limit(), Request->loop_tripcount());
 337
 338   Reply->set_number(Ret);
 339
 340   SERVER_DBG("Ran TargetTeamRegion on device %d with %d args",
 341              mapHostRTLDeviceId(Request->device_id()), Request->arg_num())
 342   return Status::OK;
 343 }
 344
 345 int32_t RemoteOffloadImpl::mapHostRTLDeviceId(int32_t RTLDeviceID) {
 346   for (auto &RTL : PM->RTLs.UsedRTLs) {
 347     if (RTLDeviceID - RTL->NumberOfDevices >= 0)
 348       RTLDeviceID -= RTL->NumberOfDevices;
 349     else
 350       break;
 351   }
 352   return RTLDeviceID;
 353 }