From 76514c889475461aa31f6bdabe3208bb49ccb19a Mon Sep 17 00:00:00 2001 From: Lang Hames Date: Tue, 23 Dec 2025 16:56:15 +1100 Subject: [PATCH] [ORC] Pass JITDispatchHandler argument buffers as WrapperFunctionBuffer. Updates ExecutionSession::runJITDispatchHandler to take the argument buffer for the function as a WrapperFunctionBuffer, rather than an ArrayRef. This is a first step towards more efficient jit-dispatch handler calls: 1. Handlers can now be run as tasks, since they own their argument buffer (so there's no risk of it being deallocated before they're run) 2. In in-process JIT setups, this will allow argument buffers to be passed in directly from the ORC runtime, rather than having to copy the buffer. --- llvm/include/llvm/ExecutionEngine/Orc/Core.h | 2 +- .../Orc/Shared/SimpleRemoteEPCUtils.h | 5 ++--- .../llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h | 10 +++++----- .../Orc/TargetProcess/SimpleRemoteEPCServer.h | 6 +++--- llvm/lib/ExecutionEngine/Orc/Core.cpp | 8 ++++---- llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp | 6 +++--- .../Orc/SelfExecutorProcessControl.cpp | 3 ++- .../Orc/Shared/SimpleRemoteEPCUtils.cpp | 7 ++++--- llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp | 14 +++++++------- .../Orc/TargetProcess/SimpleRemoteEPCServer.cpp | 6 +++--- .../ExecutionSessionWrapperFunctionCallsTest.cpp | 2 +- 11 files changed, 35 insertions(+), 34 deletions(-) diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h index 4ca63857ab464..ecdb246aad311 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h +++ b/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 ArgBuffer); + shared::WrapperFunctionBuffer ArgBytes); /// Dump the state of all the JITDylibs in this session. LLVM_ABI void dump(raw_ostream &OS); diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h index 14f16838a0154..a81ffe975b876 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h +++ b/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 BootstrapSymbols; }; -using SimpleRemoteEPCArgBytesVector = SmallVector; - class LLVM_ABI SimpleRemoteEPCTransportClient { public: enum HandleMessageAction { ContinueSession, EndSession }; @@ -65,7 +64,7 @@ class LLVM_ABI SimpleRemoteEPCTransportClient { /// otherwise. virtual Expected 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. diff --git a/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h b/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h index 81761831bdd96..83f4908ce6a98 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h @@ -86,7 +86,7 @@ class LLVM_ABI SimpleRemoteEPC : public ExecutorProcessControl, Expected 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 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) {} diff --git a/llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h b/llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h index 70a226504e876..e5d28345bb84e 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h +++ b/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 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 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); diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp index 55259f2e2bcad..121ef2eeec42f 100644 --- a/llvm/lib/ExecutionEngine/Orc/Core.cpp +++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp @@ -1897,9 +1897,9 @@ Error ExecutionSession::registerJITDispatchHandlers( return Error::success(); } -void ExecutionSession::runJITDispatchHandler(SendResultFunction SendResult, - ExecutorAddr HandlerFnTagAddr, - ArrayRef ArgBuffer) { +void ExecutionSession::runJITDispatchHandler( + SendResultFunction SendResult, ExecutorAddr HandlerFnTagAddr, + shared::WrapperFunctionBuffer ArgBytes) { std::shared_ptr 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 " + diff --git a/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp index 124279b6658d1..1e0ecc25b1971 100644 --- a/llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp +++ b/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; @@ -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"; diff --git a/llvm/lib/ExecutionEngine/Orc/SelfExecutorProcessControl.cpp b/llvm/lib/ExecutionEngine/Orc/SelfExecutorProcessControl.cpp index 8a067c9d5d165..8d49278b3a0f8 100644 --- a/llvm/lib/ExecutionEngine/Orc/SelfExecutorProcessControl.cpp +++ b/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(); } diff --git a/llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp b/llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp index e4ab2c0fd03ae..08fcd1c45445c 100644 --- a/llvm/lib/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.cpp +++ b/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 { diff --git a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp index 772aa99d8ea96..c7f8e5fb74134 100644 --- a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp +++ b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp @@ -133,7 +133,7 @@ Error SimpleRemoteEPC::disconnect() { Expected 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("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,10 +426,10 @@ 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 = @@ -437,12 +437,12 @@ void SimpleRemoteEPC::handleCallWrapper( 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()) diff --git a/llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.cpp b/llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.cpp index 781da7b319db0..5aa61526be848 100644 --- a/llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.cpp +++ b/llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.cpp @@ -62,7 +62,7 @@ StringMap SimpleRemoteEPCServer::defaultBootstrapSymbols() { Expected 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 *P = nullptr; { std::lock_guard 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); diff --git a/llvm/unittests/ExecutionEngine/Orc/ExecutionSessionWrapperFunctionCallsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ExecutionSessionWrapperFunctionCallsTest.cpp index 6c53336b49365..faab7f2df108b 100644 --- a/llvm/unittests/ExecutionEngine/Orc/ExecutionSessionWrapperFunctionCallsTest.cpp +++ b/llvm/unittests/ExecutionEngine/Orc/ExecutionSessionWrapperFunctionCallsTest.cpp @@ -110,7 +110,7 @@ TEST(ExecutionSessionWrapperFunctionCalls, RegisterAsyncHandlerAndRun) { EXPECT_TRUE(SPSArgList::deserialize(IB, Result)); RP.set_value(Result); }, - AddAsyncTagAddr, ArrayRef(ArgBuffer)); + AddAsyncTagAddr, std::move(ArgBuffer)); EXPECT_EQ(RF.get(), (int32_t)3);