llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.cpp

   1 //===------- SimpleEPCServer.cpp - EPC over simple abstract channel -------===//
   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 #include "llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h"
  10
  11 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
  12 #include "llvm/Support/FormatVariadic.h"
  13 #include "llvm/Support/Host.h"
  14 #include "llvm/Support/Process.h"
  15
  16 #include "OrcRTBootstrap.h"
  17
  18 #define DEBUG_TYPE "orc"
  19
  20 using namespace llvm::orc::shared;
  21
  22 namespace llvm {
  23 namespace orc {
  24
  25 ExecutorBootstrapService::~ExecutorBootstrapService() {}
  26
  27 SimpleRemoteEPCServer::Dispatcher::~Dispatcher() {}
  28
  29 #if LLVM_ENABLE_THREADS
  30 void SimpleRemoteEPCServer::ThreadDispatcher::dispatch(
  31     unique_function<void()> Work) {
  32   {
  33     std::lock_guard<std::mutex> Lock(DispatchMutex);
  34     if (!Running)
  35       return;
  36     ++Outstanding;
  37   }
  38
  39   std::thread([this, Work = std::move(Work)]() mutable {
  40     Work();
  41     std::lock_guard<std::mutex> Lock(DispatchMutex);
  42     --Outstanding;
  43     OutstandingCV.notify_all();
  44   }).detach();
  45 }
  46
  47 void SimpleRemoteEPCServer::ThreadDispatcher::shutdown() {
  48   std::unique_lock<std::mutex> Lock(DispatchMutex);
  49   Running = false;
  50   OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
  51 }
  52 #endif
  53
  54 StringMap<ExecutorAddr> SimpleRemoteEPCServer::defaultBootstrapSymbols() {
  55   StringMap<ExecutorAddr> DBS;
  56   rt_bootstrap::addTo(DBS);
  57   return DBS;
  58 }
  59
  60 Expected<SimpleRemoteEPCTransportClient::HandleMessageAction>
  61 SimpleRemoteEPCServer::handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
  62                                      ExecutorAddr TagAddr,
  63                                      SimpleRemoteEPCArgBytesVector ArgBytes) {
  64
  65   LLVM_DEBUG({
  66     dbgs() << "SimpleRemoteEPCServer::handleMessage: opc = ";
  67     switch (OpC) {
  68     case SimpleRemoteEPCOpcode::Setup:
  69       dbgs() << "Setup";
  70       assert(SeqNo == 0 && "Non-zero SeqNo for Setup?");
  71       assert(TagAddr.getValue() == 0 && "Non-zero TagAddr for Setup?");
  72       break;
  73     case SimpleRemoteEPCOpcode::Hangup:
  74       dbgs() << "Hangup";
  75       assert(SeqNo == 0 && "Non-zero SeqNo for Hangup?");
  76       assert(TagAddr.getValue() == 0 && "Non-zero TagAddr for Hangup?");
  77       break;
  78     case SimpleRemoteEPCOpcode::Result:
  79       dbgs() << "Result";
  80       assert(TagAddr.getValue() == 0 && "Non-zero TagAddr for Result?");
  81       break;
  82     case SimpleRemoteEPCOpcode::CallWrapper:
  83       dbgs() << "CallWrapper";
  84       break;
  85     }
  86     dbgs() << ", seqno = " << SeqNo
  87            << ", tag-addr = " << formatv("{0:x}", TagAddr.getValue())
  88            << ", arg-buffer = " << formatv("{0:x}", ArgBytes.size())
  89            << " bytes\n";
  90   });
  91
  92   using UT = std::underlying_type_t<SimpleRemoteEPCOpcode>;
  93   if (static_cast<UT>(OpC) > static_cast<UT>(SimpleRemoteEPCOpcode::LastOpC))
  94     return make_error<StringError>("Unexpected opcode",
  95                                    inconvertibleErrorCode());
  96
  97   // TODO: Clean detach message?
  98   switch (OpC) {
  99   case SimpleRemoteEPCOpcode::Setup:
 100     return make_error<StringError>("Unexpected Setup opcode",
 101                                    inconvertibleErrorCode());
 102   case SimpleRemoteEPCOpcode::Hangup:
 103     return SimpleRemoteEPCTransportClient::EndSession;
 104   case SimpleRemoteEPCOpcode::Result:
 105     if (auto Err = handleResult(SeqNo, TagAddr, std::move(ArgBytes)))
 106       return std::move(Err);
 107     break;
 108   case SimpleRemoteEPCOpcode::CallWrapper:
 109     handleCallWrapper(SeqNo, TagAddr, std::move(ArgBytes));
 110     break;
 111   }
 112   return ContinueSession;
 113 }
 114
 115 Error SimpleRemoteEPCServer::waitForDisconnect() {
 116   std::unique_lock<std::mutex> Lock(ServerStateMutex);
 117   ShutdownCV.wait(Lock, [this]() { return RunState == ServerShutDown; });
 118   return std::move(ShutdownErr);
 119 }
 120
 121 void SimpleRemoteEPCServer::handleDisconnect(Error Err) {
 122   PendingJITDispatchResultsMap TmpPending;
 123
 124   {
 125     std::lock_guard<std::mutex> Lock(ServerStateMutex);
 126     std::swap(TmpPending, PendingJITDispatchResults);
 127     RunState = ServerShuttingDown;
 128   }
 129
 130   // Send out-of-band errors to any waiting threads.
 131   for (auto &KV : TmpPending)
 132     KV.second->set_value(
 133         shared::WrapperFunctionResult::createOutOfBandError("disconnecting"));
 134
 135   // Wait for dispatcher to clear.
 136   D->shutdown();
 137
 138   // Shut down services.
 139   while (!Services.empty()) {
 140     ShutdownErr =
 141       joinErrors(std::move(ShutdownErr), Services.back()->shutdown());
 142     Services.pop_back();
 143   }
 144
 145   std::lock_guard<std::mutex> Lock(ServerStateMutex);
 146   ShutdownErr = joinErrors(std::move(ShutdownErr), std::move(Err));
 147   RunState = ServerShutDown;
 148   ShutdownCV.notify_all();
 149 }
 150
 151 Error SimpleRemoteEPCServer::sendMessage(SimpleRemoteEPCOpcode OpC,
 152                                          uint64_t SeqNo, ExecutorAddr TagAddr,
 153                                          ArrayRef<char> ArgBytes) {
 154
 155   LLVM_DEBUG({
 156     dbgs() << "SimpleRemoteEPCServer::sendMessage: opc = ";
 157     switch (OpC) {
 158     case SimpleRemoteEPCOpcode::Setup:
 159       dbgs() << "Setup";
 160       assert(SeqNo == 0 && "Non-zero SeqNo for Setup?");
 161       assert(TagAddr.getValue() == 0 && "Non-zero TagAddr for Setup?");
 162       break;
 163     case SimpleRemoteEPCOpcode::Hangup:
 164       dbgs() << "Hangup";
 165       assert(SeqNo == 0 && "Non-zero SeqNo for Hangup?");
 166       assert(TagAddr.getValue() == 0 && "Non-zero TagAddr for Hangup?");
 167       break;
 168     case SimpleRemoteEPCOpcode::Result:
 169       dbgs() << "Result";
 170       assert(TagAddr.getValue() == 0 && "Non-zero TagAddr for Result?");
 171       break;
 172     case SimpleRemoteEPCOpcode::CallWrapper:
 173       dbgs() << "CallWrapper";
 174       break;
 175     }
 176     dbgs() << ", seqno = " << SeqNo
 177            << ", tag-addr = " << formatv("{0:x}", TagAddr.getValue())
 178            << ", arg-buffer = " << formatv("{0:x}", ArgBytes.size())
 179            << " bytes\n";
 180   });
 181   auto Err = T->sendMessage(OpC, SeqNo, TagAddr, ArgBytes);
 182   LLVM_DEBUG({
 183     if (Err)
 184       dbgs() << "  \\--> SimpleRemoteEPC::sendMessage failed\n";
 185   });
 186   return Err;
 187 }
 188
 189 Error SimpleRemoteEPCServer::sendSetupMessage(
 190     StringMap<ExecutorAddr> BootstrapSymbols) {
 191
 192   using namespace SimpleRemoteEPCDefaultBootstrapSymbolNames;
 193
 194   std::vector<char> SetupPacket;
 195   SimpleRemoteEPCExecutorInfo EI;
 196   EI.TargetTriple = sys::getProcessTriple();
 197   if (auto PageSize = sys::Process::getPageSize())
 198     EI.PageSize = *PageSize;
 199   else
 200     return PageSize.takeError();
 201   EI.BootstrapSymbols = std::move(BootstrapSymbols);
 202
 203   assert(!EI.BootstrapSymbols.count(ExecutorSessionObjectName) &&
 204          "Dispatch context name should not be set");
 205   assert(!EI.BootstrapSymbols.count(DispatchFnName) &&
 206          "Dispatch function name should not be set");
 207   EI.BootstrapSymbols[ExecutorSessionObjectName] = ExecutorAddr::fromPtr(this);
 208   EI.BootstrapSymbols[DispatchFnName] = ExecutorAddr::fromPtr(jitDispatchEntry);
 209
 210   using SPSSerialize =
 211       shared::SPSArgList<shared::SPSSimpleRemoteEPCExecutorInfo>;
 212   auto SetupPacketBytes =
 213       shared::WrapperFunctionResult::allocate(SPSSerialize::size(EI));
 214   shared::SPSOutputBuffer OB(SetupPacketBytes.data(), SetupPacketBytes.size());
 215   if (!SPSSerialize::serialize(OB, EI))
 216     return make_error<StringError>("Could not send setup packet",
 217                                    inconvertibleErrorCode());
 218
 219   return sendMessage(SimpleRemoteEPCOpcode::Setup, 0, ExecutorAddr(),
 220                      {SetupPacketBytes.data(), SetupPacketBytes.size()});
 221 }
 222
 223 Error SimpleRemoteEPCServer::handleResult(
 224     uint64_t SeqNo, ExecutorAddr TagAddr,
 225     SimpleRemoteEPCArgBytesVector ArgBytes) {
 226   std::promise<shared::WrapperFunctionResult> *P = nullptr;
 227   {
 228     std::lock_guard<std::mutex> Lock(ServerStateMutex);
 229     auto I = PendingJITDispatchResults.find(SeqNo);
 230     if (I == PendingJITDispatchResults.end())
 231       return make_error<StringError>("No call for sequence number " +
 232                                          Twine(SeqNo),
 233                                      inconvertibleErrorCode());
 234     P = I->second;
 235     PendingJITDispatchResults.erase(I);
 236     releaseSeqNo(SeqNo);
 237   }
 238   auto R = shared::WrapperFunctionResult::allocate(ArgBytes.size());
 239   memcpy(R.data(), ArgBytes.data(), ArgBytes.size());
 240   P->set_value(std::move(R));
 241   return Error::success();
 242 }
 243
 244 void SimpleRemoteEPCServer::handleCallWrapper(
 245     uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
 246     SimpleRemoteEPCArgBytesVector ArgBytes) {
 247   D->dispatch([this, RemoteSeqNo, TagAddr, ArgBytes = std::move(ArgBytes)]() {
 248     using WrapperFnTy =
 249         shared::CWrapperFunctionResult (*)(const char *, size_t);
 250     auto *Fn = TagAddr.toPtr<WrapperFnTy>();
 251     shared::WrapperFunctionResult ResultBytes(
 252         Fn(ArgBytes.data(), ArgBytes.size()));
 253     if (auto Err = sendMessage(SimpleRemoteEPCOpcode::Result, RemoteSeqNo,
 254                                ExecutorAddr(),
 255                                {ResultBytes.data(), ResultBytes.size()}))
 256       ReportError(std::move(Err));
 257   });
 258 }
 259
 260 shared::WrapperFunctionResult
 261 SimpleRemoteEPCServer::doJITDispatch(const void *FnTag, const char *ArgData,
 262                                      size_t ArgSize) {
 263   uint64_t SeqNo;
 264   std::promise<shared::WrapperFunctionResult> ResultP;
 265   auto ResultF = ResultP.get_future();
 266   {
 267     std::lock_guard<std::mutex> Lock(ServerStateMutex);
 268     if (RunState != ServerRunning)
 269       return shared::WrapperFunctionResult::createOutOfBandError(
 270           "jit_dispatch not available (EPC server shut down)");
 271
 272     SeqNo = getNextSeqNo();
 273     assert(!PendingJITDispatchResults.count(SeqNo) && "SeqNo already in use");
 274     PendingJITDispatchResults[SeqNo] = &ResultP;
 275   }
 276
 277   if (auto Err = sendMessage(SimpleRemoteEPCOpcode::CallWrapper, SeqNo,
 278                              ExecutorAddr::fromPtr(FnTag), {ArgData, ArgSize}))
 279     ReportError(std::move(Err));
 280
 281   return ResultF.get();
 282 }
 283
 284 shared::CWrapperFunctionResult
 285 SimpleRemoteEPCServer::jitDispatchEntry(void *DispatchCtx, const void *FnTag,
 286                                         const char *ArgData, size_t ArgSize) {
 287   return reinterpret_cast<SimpleRemoteEPCServer *>(DispatchCtx)
 288       ->doJITDispatch(FnTag, ArgData, ArgSize)
 289       .release();
 290 }
 291
 292 } // end namespace orc
 293 } // end namespace llvm