[ORC] Pass JITDispatchHandler argument buffers as WrapperFunctionBuffer.

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -1692,7 +1692,7 @@ class ExecutionSession {
   /// to incoming jit-dispatch requests from the executor.
   LLVM_ABI void runJITDispatchHandler(SendResultFunction SendResult,
                                       ExecutorAddr HandlerFnTagAddr,
-                                      ArrayRef<char> ArgBuffer);
+                                      shared::WrapperFunctionBuffer ArgBytes);
 
   /// Dump the state of all the JITDylibs in this session.
   LLVM_ABI void dump(raw_ostream &OS);

llvm/include/llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h

@@ -19,6 +19,7 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
+#include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Error.h"
 
@@ -50,8 +51,6 @@ struct SimpleRemoteEPCExecutorInfo {
   StringMap<ExecutorAddr> BootstrapSymbols;
 };
 
-using SimpleRemoteEPCArgBytesVector = SmallVector<char, 128>;
-
 class LLVM_ABI SimpleRemoteEPCTransportClient {
 public:
   enum HandleMessageAction { ContinueSession, EndSession };
@@ -65,7 +64,7 @@ class LLVM_ABI SimpleRemoteEPCTransportClient {
   /// otherwise.
   virtual Expected<HandleMessageAction>
   handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr,
-                SimpleRemoteEPCArgBytesVector ArgBytes) = 0;
+                shared::WrapperFunctionBuffer ArgBytes) = 0;
 
   /// Handle a disconnection from the underlying transport. No further messages
   /// should be sent to handleMessage after this is called.

llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h

@@ -86,7 +86,7 @@ class LLVM_ABI SimpleRemoteEPC : public ExecutorProcessControl,
 
   Expected<HandleMessageAction>
   handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr,
-                SimpleRemoteEPCArgBytesVector ArgBytes) override;
+                shared::WrapperFunctionBuffer ArgBytes) override;
 
   void handleDisconnect(Error Err) override;
 
@@ -106,14 +106,14 @@ class LLVM_ABI SimpleRemoteEPC : public ExecutorProcessControl,
                     ExecutorAddr TagAddr, ArrayRef<char> ArgBytes);
 
   Error handleSetup(uint64_t SeqNo, ExecutorAddr TagAddr,
-                    SimpleRemoteEPCArgBytesVector ArgBytes);
+                    shared::WrapperFunctionBuffer ArgBytes);
   Error setup(Setup S);
 
   Error handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
-                     SimpleRemoteEPCArgBytesVector ArgBytes);
+                     shared::WrapperFunctionBuffer ArgBytes);
   void handleCallWrapper(uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
-                         SimpleRemoteEPCArgBytesVector ArgBytes);
-  Error handleHangup(SimpleRemoteEPCArgBytesVector ArgBytes);
+                         shared::WrapperFunctionBuffer ArgBytes);
+  Error handleHangup(shared::WrapperFunctionBuffer ArgBytes);
 
   uint64_t getNextSeqNo() { return NextSeqNo++; }
   void releaseSeqNo(uint64_t SeqNo) {}

llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h

@@ -145,7 +145,7 @@ class LLVM_ABI SimpleRemoteEPCServer : public SimpleRemoteEPCTransportClient {
   /// returns an error, which should be reported and treated as a 'Disconnect'.
   Expected<HandleMessageAction>
   handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr,
-                SimpleRemoteEPCArgBytesVector ArgBytes) override;
+                shared::WrapperFunctionBuffer ArgBytes) override;
 
   Error waitForDisconnect();
 
@@ -159,9 +159,9 @@ class LLVM_ABI SimpleRemoteEPCServer : public SimpleRemoteEPCTransportClient {
                          StringMap<ExecutorAddr> BootstrapSymbols);
 
   Error handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
-                     SimpleRemoteEPCArgBytesVector ArgBytes);
+                     shared::WrapperFunctionBuffer ArgBytes);
   void handleCallWrapper(uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
-                         SimpleRemoteEPCArgBytesVector ArgBytes);
+                         shared::WrapperFunctionBuffer ArgBytes);
 
   shared::WrapperFunctionBuffer
   doJITDispatch(const void *FnTag, const char *ArgData, size_t ArgSize);

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -1897,9 +1897,9 @@ Error ExecutionSession::registerJITDispatchHandlers(
   return Error::success();
 }
 
-void ExecutionSession::runJITDispatchHandler(SendResultFunction SendResult,
-                                             ExecutorAddr HandlerFnTagAddr,
-                                             ArrayRef<char> ArgBuffer) {
+void ExecutionSession::runJITDispatchHandler(
+    SendResultFunction SendResult, ExecutorAddr HandlerFnTagAddr,
+    shared::WrapperFunctionBuffer ArgBytes) {
 
   std::shared_ptr<JITDispatchHandlerFunction> F;
   {
@@ -1910,7 +1910,7 @@ void ExecutionSession::runJITDispatchHandler(SendResultFunction SendResult,
   }
 
   if (F)
-    (*F)(std::move(SendResult), ArgBuffer.data(), ArgBuffer.size());
+    (*F)(std::move(SendResult), ArgBytes.data(), ArgBytes.size());
   else
     SendResult(shared::WrapperFunctionBuffer::createOutOfBandError(
         ("No function registered for tag " +

llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp

@@ -27,8 +27,8 @@ void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
 }
 
 static void orc_rt_lite_reoptimize_helper(
-    shared::CWrapperFunctionBuffer (*JITDispatch)(void *Ctx, void *Tag,
-                                                  void *Data, size_t Size),
+    shared::CWrapperFunctionBuffer (*JITDispatch)(
+        void *Ctx, void *Tag, shared::CWrapperFunctionBuffer),
     void *JITDispatchCtx, void *Tag, uint64_t MUID, uint32_t CurVersion) {
   // Serialize the arguments into a WrapperFunctionBuffer and call dispatch.
   using SPSArgs = shared::SPSArgList<uint64_t, uint32_t>;
@@ -41,7 +41,7 @@ static void orc_rt_lite_reoptimize_helper(
     abort();
   }
   shared::WrapperFunctionBuffer Buf{
-      JITDispatch(JITDispatchCtx, Tag, ArgBytes.data(), ArgBytes.size())};
+      JITDispatch(JITDispatchCtx, Tag, ArgBytes.release())};
 
   if (const char *ErrMsg = Buf.getOutOfBandError()) {
     errs() << "Reoptimization error: " << ErrMsg << "\naborting.\n";

llvm/lib/ExecutionEngine/Orc/SelfExecutorProcessControl.cpp

@@ -161,7 +161,8 @@ SelfExecutorProcessControl::jitDispatchViaWrapperFunctionManager(
               shared::WrapperFunctionBuffer Result) mutable {
             ResultP.set_value(std::move(Result));
           },
-          ExecutorAddr::fromPtr(FnTag), {Data, Size});
+          ExecutorAddr::fromPtr(FnTag),
+          shared::WrapperFunctionBuffer::copyFrom(Data, Size));
 
   return ResultF.get().release();
 }

llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp

@@ -222,14 +222,15 @@ void FDSimpleRemoteEPCTransport::listenLoop() {
     }
 
     // Read the argument bytes.
-    SimpleRemoteEPCArgBytesVector ArgBytes;
-    ArgBytes.resize(MsgSize - FDMsgHeader::Size);
+    auto ArgBytes =
+        shared::WrapperFunctionBuffer::allocate(MsgSize - FDMsgHeader::Size);
     if (auto Err2 = readBytes(ArgBytes.data(), ArgBytes.size())) {
       Err = joinErrors(std::move(Err), std::move(Err2));
       break;
     }
 
-    if (auto Action = C.handleMessage(OpC, SeqNo, TagAddr, ArgBytes)) {
+    if (auto Action =
+            C.handleMessage(OpC, SeqNo, TagAddr, std::move(ArgBytes))) {
       if (*Action == SimpleRemoteEPCTransportClient::EndSession)
         break;
     } else {

llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp

@@ -133,7 +133,7 @@ Error SimpleRemoteEPC::disconnect() {
 Expected<SimpleRemoteEPCTransportClient::HandleMessageAction>
 SimpleRemoteEPC::handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
                                ExecutorAddr TagAddr,
-                               SimpleRemoteEPCArgBytesVector ArgBytes) {
+                               shared::WrapperFunctionBuffer ArgBytes) {
 
   LLVM_DEBUG({
     dbgs() << "SimpleRemoteEPC::handleMessage: opc = ";
@@ -282,7 +282,7 @@ Error SimpleRemoteEPC::sendMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
 }
 
 Error SimpleRemoteEPC::handleSetup(uint64_t SeqNo, ExecutorAddr TagAddr,
-                                   SimpleRemoteEPCArgBytesVector ArgBytes) {
+                                   shared::WrapperFunctionBuffer ArgBytes) {
   if (SeqNo != 0)
     return make_error<StringError>("Setup packet SeqNo not zero",
                                    inconvertibleErrorCode());
@@ -399,7 +399,7 @@ Error SimpleRemoteEPC::setup(Setup S) {
 }
 
 Error SimpleRemoteEPC::handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
-                                    SimpleRemoteEPCArgBytesVector ArgBytes) {
+                                    shared::WrapperFunctionBuffer ArgBytes) {
   IncomingWFRHandler SendResult;
 
   if (TagAddr)
@@ -426,23 +426,23 @@ Error SimpleRemoteEPC::handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
 
 void SimpleRemoteEPC::handleCallWrapper(
     uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
-    SimpleRemoteEPCArgBytesVector ArgBytes) {
+    shared::WrapperFunctionBuffer ArgBytes) {
   assert(ES && "No ExecutionSession attached");
   D->dispatch(makeGenericNamedTask(
-      [this, RemoteSeqNo, TagAddr, ArgBytes = std::move(ArgBytes)]() {
+      [this, RemoteSeqNo, TagAddr, ArgBytes = std::move(ArgBytes)]() mutable {
         ES->runJITDispatchHandler(
             [this, RemoteSeqNo](shared::WrapperFunctionBuffer WFR) {
               if (auto Err =
                       sendMessage(SimpleRemoteEPCOpcode::Result, RemoteSeqNo,
                                   ExecutorAddr(), {WFR.data(), WFR.size()}))
                 getExecutionSession().reportError(std::move(Err));
             },
-            TagAddr, ArgBytes);
+            TagAddr, std::move(ArgBytes));
       },
       "callWrapper task"));
 }
 
-Error SimpleRemoteEPC::handleHangup(SimpleRemoteEPCArgBytesVector ArgBytes) {
+Error SimpleRemoteEPC::handleHangup(shared::WrapperFunctionBuffer ArgBytes) {
   using namespace llvm::orc::shared;
   auto WFR = WrapperFunctionBuffer::copyFrom(ArgBytes.data(), ArgBytes.size());
   if (const char *ErrMsg = WFR.getOutOfBandError())

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

@@ -62,7 +62,7 @@ StringMap<ExecutorAddr> SimpleRemoteEPCServer::defaultBootstrapSymbols() {
 Expected<SimpleRemoteEPCTransportClient::HandleMessageAction>
 SimpleRemoteEPCServer::handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
                                      ExecutorAddr TagAddr,
-                                     SimpleRemoteEPCArgBytesVector ArgBytes) {
+                                     shared::WrapperFunctionBuffer ArgBytes) {
 
   LLVM_DEBUG({
     dbgs() << "SimpleRemoteEPCServer::handleMessage: opc = ";
@@ -224,7 +224,7 @@ Error SimpleRemoteEPCServer::sendSetupMessage(
 
 Error SimpleRemoteEPCServer::handleResult(
     uint64_t SeqNo, ExecutorAddr TagAddr,
-    SimpleRemoteEPCArgBytesVector ArgBytes) {
+    shared::WrapperFunctionBuffer ArgBytes) {
   std::promise<shared::WrapperFunctionBuffer> *P = nullptr;
   {
     std::lock_guard<std::mutex> Lock(ServerStateMutex);
@@ -245,7 +245,7 @@ Error SimpleRemoteEPCServer::handleResult(
 
 void SimpleRemoteEPCServer::handleCallWrapper(
     uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
-    SimpleRemoteEPCArgBytesVector ArgBytes) {
+    shared::WrapperFunctionBuffer ArgBytes) {
   D->dispatch([this, RemoteSeqNo, TagAddr, ArgBytes = std::move(ArgBytes)]() {
     using WrapperFnTy =
         shared::CWrapperFunctionBuffer (*)(const char *, size_t);

llvm/unittests/ExecutionEngine/Orc/ExecutionSessionWrapperFunctionCallsTest.cpp

@@ -110,7 +110,7 @@ TEST(ExecutionSessionWrapperFunctionCalls, RegisterAsyncHandlerAndRun) {
         EXPECT_TRUE(SPSArgList<int32_t>::deserialize(IB, Result));
         RP.set_value(Result);
       },
-      AddAsyncTagAddr, ArrayRef<char>(ArgBuffer));
+      AddAsyncTagAddr, std::move(ArgBuffer));
 
   EXPECT_EQ(RF.get(), (int32_t)3);