Skip to content

Commit fa7f7a4

Browse files
lhameslhames
authored
[ORC] Pass JITDispatchHandler argument buffers as WrapperFunctionBuffer. (#173334)
Updates ExecutionSession::runJITDispatchHandler to take the argument buffer for the function as a WrapperFunctionBuffer, rather than an ArrayRef<char>. 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.
1 parent f4ded95 commit fa7f7a4

File tree

11 files changed

+35
-34
lines changed

11 files changed

+35
-34
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1692,7 +1692,7 @@ class ExecutionSession {
16921692
/// to incoming jit-dispatch requests from the executor.
16931693
LLVM_ABI void runJITDispatchHandler(SendResultFunction SendResult,
16941694
ExecutorAddr HandlerFnTagAddr,
1695-
ArrayRef<char> ArgBuffer);
1695+
shared::WrapperFunctionBuffer ArgBytes);
16961696

16971697
/// Dump the state of all the JITDylibs in this session.
16981698
LLVM_ABI void dump(raw_ostream &OS);

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "llvm/ADT/StringMap.h"
2020
#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
2121
#include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
22+
#include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
2223
#include "llvm/Support/Compiler.h"
2324
#include "llvm/Support/Error.h"
2425

@@ -50,8 +51,6 @@ struct SimpleRemoteEPCExecutorInfo {
5051
StringMap<ExecutorAddr> BootstrapSymbols;
5152
};
5253

53-
using SimpleRemoteEPCArgBytesVector = SmallVector<char, 128>;
54-
5554
class LLVM_ABI SimpleRemoteEPCTransportClient {
5655
public:
5756
enum HandleMessageAction { ContinueSession, EndSession };
@@ -65,7 +64,7 @@ class LLVM_ABI SimpleRemoteEPCTransportClient {
6564
/// otherwise.
6665
virtual Expected<HandleMessageAction>
6766
handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr,
68-
SimpleRemoteEPCArgBytesVector ArgBytes) = 0;
67+
shared::WrapperFunctionBuffer ArgBytes) = 0;
6968

7069
/// Handle a disconnection from the underlying transport. No further messages
7170
/// should be sent to handleMessage after this is called.

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ class LLVM_ABI SimpleRemoteEPC : public ExecutorProcessControl,
8686

8787
Expected<HandleMessageAction>
8888
handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr,
89-
SimpleRemoteEPCArgBytesVector ArgBytes) override;
89+
shared::WrapperFunctionBuffer ArgBytes) override;
9090

9191
void handleDisconnect(Error Err) override;
9292

@@ -106,14 +106,14 @@ class LLVM_ABI SimpleRemoteEPC : public ExecutorProcessControl,
106106
ExecutorAddr TagAddr, ArrayRef<char> ArgBytes);
107107

108108
Error handleSetup(uint64_t SeqNo, ExecutorAddr TagAddr,
109-
SimpleRemoteEPCArgBytesVector ArgBytes);
109+
shared::WrapperFunctionBuffer ArgBytes);
110110
Error setup(Setup S);
111111

112112
Error handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
113-
SimpleRemoteEPCArgBytesVector ArgBytes);
113+
shared::WrapperFunctionBuffer ArgBytes);
114114
void handleCallWrapper(uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
115-
SimpleRemoteEPCArgBytesVector ArgBytes);
116-
Error handleHangup(SimpleRemoteEPCArgBytesVector ArgBytes);
115+
shared::WrapperFunctionBuffer ArgBytes);
116+
Error handleHangup(shared::WrapperFunctionBuffer ArgBytes);
117117

118118
uint64_t getNextSeqNo() { return NextSeqNo++; }
119119
void releaseSeqNo(uint64_t SeqNo) {}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ class LLVM_ABI SimpleRemoteEPCServer : public SimpleRemoteEPCTransportClient {
145145
/// returns an error, which should be reported and treated as a 'Disconnect'.
146146
Expected<HandleMessageAction>
147147
handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo, ExecutorAddr TagAddr,
148-
SimpleRemoteEPCArgBytesVector ArgBytes) override;
148+
shared::WrapperFunctionBuffer ArgBytes) override;
149149

150150
Error waitForDisconnect();
151151

@@ -159,9 +159,9 @@ class LLVM_ABI SimpleRemoteEPCServer : public SimpleRemoteEPCTransportClient {
159159
StringMap<ExecutorAddr> BootstrapSymbols);
160160

161161
Error handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
162-
SimpleRemoteEPCArgBytesVector ArgBytes);
162+
shared::WrapperFunctionBuffer ArgBytes);
163163
void handleCallWrapper(uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
164-
SimpleRemoteEPCArgBytesVector ArgBytes);
164+
shared::WrapperFunctionBuffer ArgBytes);
165165

166166
shared::WrapperFunctionBuffer
167167
doJITDispatch(const void *FnTag, const char *ArgData, size_t ArgSize);

‎llvm/lib/ExecutionEngine/Orc/Core.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1897,9 +1897,9 @@ Error ExecutionSession::registerJITDispatchHandlers(
18971897
return Error::success();
18981898
}
18991899

1900-
void ExecutionSession::runJITDispatchHandler(SendResultFunction SendResult,
1901-
ExecutorAddr HandlerFnTagAddr,
1902-
ArrayRef<char> ArgBuffer) {
1900+
void ExecutionSession::runJITDispatchHandler(
1901+
SendResultFunction SendResult, ExecutorAddr HandlerFnTagAddr,
1902+
shared::WrapperFunctionBuffer ArgBytes) {
19031903

19041904
std::shared_ptr<JITDispatchHandlerFunction> F;
19051905
{
@@ -1910,7 +1910,7 @@ void ExecutionSession::runJITDispatchHandler(SendResultFunction SendResult,
19101910
}
19111911

19121912
if (F)
1913-
(*F)(std::move(SendResult), ArgBuffer.data(), ArgBuffer.size());
1913+
(*F)(std::move(SendResult), ArgBytes.data(), ArgBytes.size());
19141914
else
19151915
SendResult(shared::WrapperFunctionBuffer::createOutOfBandError(
19161916
("No function registered for tag " +

‎llvm/lib/ExecutionEngine/Orc/ReOptimizeLayer.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ void ReOptimizeLayer::ReOptMaterializationUnitState::reoptimizeFailed() {
2727
}
2828

2929
static void orc_rt_lite_reoptimize_helper(
30-
shared::CWrapperFunctionBuffer (*JITDispatch)(void *Ctx, void *Tag,
31-
void *Data, size_t Size),
30+
shared::CWrapperFunctionBuffer (*JITDispatch)(
31+
void *Ctx, void *Tag, shared::CWrapperFunctionBuffer),
3232
void *JITDispatchCtx, void *Tag, uint64_t MUID, uint32_t CurVersion) {
3333
// Serialize the arguments into a WrapperFunctionBuffer and call dispatch.
3434
using SPSArgs = shared::SPSArgList<uint64_t, uint32_t>;
@@ -41,7 +41,7 @@ static void orc_rt_lite_reoptimize_helper(
4141
abort();
4242
}
4343
shared::WrapperFunctionBuffer Buf{
44-
JITDispatch(JITDispatchCtx, Tag, ArgBytes.data(), ArgBytes.size())};
44+
JITDispatch(JITDispatchCtx, Tag, ArgBytes.release())};
4545

4646
if (const char *ErrMsg = Buf.getOutOfBandError()) {
4747
errs() << "Reoptimization error: " << ErrMsg << "\naborting.\n";

‎llvm/lib/ExecutionEngine/Orc/SelfExecutorProcessControl.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,8 @@ SelfExecutorProcessControl::jitDispatchViaWrapperFunctionManager(
161161
shared::WrapperFunctionBuffer Result) mutable {
162162
ResultP.set_value(std::move(Result));
163163
},
164-
ExecutorAddr::fromPtr(FnTag), {Data, Size});
164+
ExecutorAddr::fromPtr(FnTag),
165+
shared::WrapperFunctionBuffer::copyFrom(Data, Size));
165166

166167
return ResultF.get().release();
167168
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -222,14 +222,15 @@ void FDSimpleRemoteEPCTransport::listenLoop() {
222222
}
223223

224224
// Read the argument bytes.
225-
SimpleRemoteEPCArgBytesVector ArgBytes;
226-
ArgBytes.resize(MsgSize - FDMsgHeader::Size);
225+
auto ArgBytes =
226+
shared::WrapperFunctionBuffer::allocate(MsgSize - FDMsgHeader::Size);
227227
if (auto Err2 = readBytes(ArgBytes.data(), ArgBytes.size())) {
228228
Err = joinErrors(std::move(Err), std::move(Err2));
229229
break;
230230
}
231231

232-
if (auto Action = C.handleMessage(OpC, SeqNo, TagAddr, ArgBytes)) {
232+
if (auto Action =
233+
C.handleMessage(OpC, SeqNo, TagAddr, std::move(ArgBytes))) {
233234
if (*Action == SimpleRemoteEPCTransportClient::EndSession)
234235
break;
235236
} else {

‎llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ Error SimpleRemoteEPC::disconnect() {
133133
Expected<SimpleRemoteEPCTransportClient::HandleMessageAction>
134134
SimpleRemoteEPC::handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
135135
ExecutorAddr TagAddr,
136-
SimpleRemoteEPCArgBytesVector ArgBytes) {
136+
shared::WrapperFunctionBuffer ArgBytes) {
137137

138138
LLVM_DEBUG({
139139
dbgs() << "SimpleRemoteEPC::handleMessage: opc = ";
@@ -282,7 +282,7 @@ Error SimpleRemoteEPC::sendMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
282282
}
283283

284284
Error SimpleRemoteEPC::handleSetup(uint64_t SeqNo, ExecutorAddr TagAddr,
285-
SimpleRemoteEPCArgBytesVector ArgBytes) {
285+
shared::WrapperFunctionBuffer ArgBytes) {
286286
if (SeqNo != 0)
287287
return make_error<StringError>("Setup packet SeqNo not zero",
288288
inconvertibleErrorCode());
@@ -399,7 +399,7 @@ Error SimpleRemoteEPC::setup(Setup S) {
399399
}
400400

401401
Error SimpleRemoteEPC::handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
402-
SimpleRemoteEPCArgBytesVector ArgBytes) {
402+
shared::WrapperFunctionBuffer ArgBytes) {
403403
IncomingWFRHandler SendResult;
404404

405405
if (TagAddr)
@@ -426,23 +426,23 @@ Error SimpleRemoteEPC::handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
426426

427427
void SimpleRemoteEPC::handleCallWrapper(
428428
uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
429-
SimpleRemoteEPCArgBytesVector ArgBytes) {
429+
shared::WrapperFunctionBuffer ArgBytes) {
430430
assert(ES && "No ExecutionSession attached");
431431
D->dispatch(makeGenericNamedTask(
432-
[this, RemoteSeqNo, TagAddr, ArgBytes = std::move(ArgBytes)]() {
432+
[this, RemoteSeqNo, TagAddr, ArgBytes = std::move(ArgBytes)]() mutable {
433433
ES->runJITDispatchHandler(
434434
[this, RemoteSeqNo](shared::WrapperFunctionBuffer WFR) {
435435
if (auto Err =
436436
sendMessage(SimpleRemoteEPCOpcode::Result, RemoteSeqNo,
437437
ExecutorAddr(), {WFR.data(), WFR.size()}))
438438
getExecutionSession().reportError(std::move(Err));
439439
},
440-
TagAddr, ArgBytes);
440+
TagAddr, std::move(ArgBytes));
441441
},
442442
"callWrapper task"));
443443
}
444444

445-
Error SimpleRemoteEPC::handleHangup(SimpleRemoteEPCArgBytesVector ArgBytes) {
445+
Error SimpleRemoteEPC::handleHangup(shared::WrapperFunctionBuffer ArgBytes) {
446446
using namespace llvm::orc::shared;
447447
auto WFR = WrapperFunctionBuffer::copyFrom(ArgBytes.data(), ArgBytes.size());
448448
if (const char *ErrMsg = WFR.getOutOfBandError())

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ StringMap<ExecutorAddr> SimpleRemoteEPCServer::defaultBootstrapSymbols() {
6262
Expected<SimpleRemoteEPCTransportClient::HandleMessageAction>
6363
SimpleRemoteEPCServer::handleMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
6464
ExecutorAddr TagAddr,
65-
SimpleRemoteEPCArgBytesVector ArgBytes) {
65+
shared::WrapperFunctionBuffer ArgBytes) {
6666

6767
LLVM_DEBUG({
6868
dbgs() << "SimpleRemoteEPCServer::handleMessage: opc = ";
@@ -224,7 +224,7 @@ Error SimpleRemoteEPCServer::sendSetupMessage(
224224

225225
Error SimpleRemoteEPCServer::handleResult(
226226
uint64_t SeqNo, ExecutorAddr TagAddr,
227-
SimpleRemoteEPCArgBytesVector ArgBytes) {
227+
shared::WrapperFunctionBuffer ArgBytes) {
228228
std::promise<shared::WrapperFunctionBuffer> *P = nullptr;
229229
{
230230
std::lock_guard<std::mutex> Lock(ServerStateMutex);
@@ -245,7 +245,7 @@ Error SimpleRemoteEPCServer::handleResult(
245245

246246
void SimpleRemoteEPCServer::handleCallWrapper(
247247
uint64_t RemoteSeqNo, ExecutorAddr TagAddr,
248-
SimpleRemoteEPCArgBytesVector ArgBytes) {
248+
shared::WrapperFunctionBuffer ArgBytes) {
249249
D->dispatch([this, RemoteSeqNo, TagAddr, ArgBytes = std::move(ArgBytes)]() {
250250
using WrapperFnTy =
251251
shared::CWrapperFunctionBuffer (*)(const char *, size_t);

0 commit comments

Comments
 (0)