Skip to content

T
tic
Unverified Commit 279a8eba by Tianqi Chen Committed by GitHub
Options

[RUNTIME][RPC] Update RPC runtime to allow remote module as arg (#4462)

parent 77bdd5f7
Showing with 91 additions and 67 deletions
python/tvm/contrib/debugger/debug_runtime.py
......@@ -23,7 +23,6 @@ from tvm._ffi.base import string_types
from tvm._ffi.function import get_global_func
from tvm.contrib import graph_runtime
from tvm.ndarray import array
from tvm.rpc import base as rpc_base
from . import debug_result
_DUMP_ROOT_PREFIX = "tvmdbg_"
......@@ -60,25 +59,17 @@ def create(graph_json_str, libmod, ctx, dump_root=None):
except AttributeError:
raise ValueError("Type %s is not supported" % type(graph_json_str))
try:
fcreate = get_global_func("tvm.graph_runtime_debug.create")
ctx, num_rpc_ctx, device_type_id = graph_runtime.get_device_ctx(libmod, ctx)
if num_rpc_ctx == len(ctx):
fcreate = ctx[0]._rpc_sess.get_function(
"tvm.graph_runtime_debug.create")
else:
fcreate = get_global_func("tvm.graph_runtime_debug.create")
except ValueError:
raise ValueError(
"Please set '(USE_GRAPH_RUNTIME_DEBUG ON)' in "
"config.cmake and rebuild TVM to enable debug mode"
)
ctx, num_rpc_ctx, device_type_id = graph_runtime.get_device_ctx(libmod, ctx)
if num_rpc_ctx == len(ctx):
libmod = rpc_base._ModuleHandle(libmod)
try:
fcreate = ctx[0]._rpc_sess.get_function(
"tvm.graph_runtime_debug.remote_create"
)
except ValueError:
raise ValueError(
"Please set '(USE_GRAPH_RUNTIME_DEBUG ON)' in "
"config.cmake and rebuild TVM to enable debug mode"
)
func_obj = fcreate(graph_json_str, libmod, *device_type_id)
return GraphModuleDebug(func_obj, ctx, graph_json_str, dump_root)
......
python/tvm/contrib/graph_runtime.py
......@@ -51,11 +51,10 @@ def create(graph_json_str, libmod, ctx):
ctx, num_rpc_ctx, device_type_id = get_device_ctx(libmod, ctx)
if num_rpc_ctx == len(ctx):
hmod = rpc_base._ModuleHandle(libmod)
fcreate = ctx[0]._rpc_sess.get_function("tvm.graph_runtime.remote_create")
return GraphModule(fcreate(graph_json_str, hmod, *device_type_id))
fcreate = ctx[0]._rpc_sess.get_function("tvm.graph_runtime.create")
else:
fcreate = get_global_func("tvm.graph_runtime.create")
fcreate = get_global_func("tvm.graph_runtime.create")
return GraphModule(fcreate(graph_json_str, libmod, *device_type_id))
def get_device_ctx(libmod, ctx):
......
src/runtime/graph/debug/graph_runtime_debug.cc
......@@ -27,7 +27,6 @@
#include <chrono>
#include <sstream>
#include "../graph_runtime.h"
#include "../../object_internal.h"
namespace tvm {
namespace runtime {
......@@ -220,19 +219,5 @@ TVM_REGISTER_GLOBAL("tvm.graph_runtime_debug.create")
<< args.num_args;
*rv = GraphRuntimeDebugCreate(args[0], args[1], GetAllContext(args));
});
TVM_REGISTER_GLOBAL("tvm.graph_runtime_debug.remote_create")
.set_body([](TVMArgs args, TVMRetValue* rv) {
CHECK_GE(args.num_args, 4) << "The expected number of arguments for "
"graph_runtime.remote_create is "
"at least 4, but it has "
<< args.num_args;
void* mhandle = args[1];
ModuleNode* mnode = ObjectInternal::GetModuleNode(mhandle);
const auto& contexts = GetAllContext(args);
*rv = GraphRuntimeDebugCreate(
args[0], GetRef<Module>(mnode), contexts);
});
} // namespace runtime
} // namespace tvm
src/runtime/graph/graph_runtime.cc
......@@ -36,7 +36,6 @@
#include <vector>
#include "graph_runtime.h"
#include "../object_internal.h"
namespace tvm {
namespace runtime {
......@@ -511,19 +510,5 @@ TVM_REGISTER_GLOBAL("tvm.graph_runtime.create")
const auto& contexts = GetAllContext(args);
*rv = GraphRuntimeCreate(args[0], args[1], contexts);
});
TVM_REGISTER_GLOBAL("tvm.graph_runtime.remote_create")
.set_body([](TVMArgs args, TVMRetValue* rv) {
CHECK_GE(args.num_args, 4) << "The expected number of arguments for "
"graph_runtime.remote_create is "
"at least 4, but it has "
<< args.num_args;
void* mhandle = args[1];
ModuleNode* mnode = ObjectInternal::GetModuleNode(mhandle);
const auto& contexts = GetAllContext(args);
*rv = GraphRuntimeCreate(
args[0], GetRef<Module>(mnode), contexts);
});
} // namespace runtime
} // namespace tvm
src/runtime/rpc/rpc_module.cc
......@@ -41,7 +41,7 @@ class RPCWrappedFunc {
}
void operator()(TVMArgs args, TVMRetValue *rv) const {
sess_->CallFunc(handle_, args, rv, &fwrap_);
sess_->CallFunc(handle_, args, rv, UnwrapRemote, &fwrap_);
}
~RPCWrappedFunc() {
try {
......@@ -55,6 +55,9 @@ class RPCWrappedFunc {
TVMArgs args,
TVMRetValue* rv);
static void* UnwrapRemote(int rpc_sess_table_index,
const TVMArgValue& arg);
// deleter of RPC remote array
static void RemoteNDArrayDeleter(NDArray::Container* ptr) {
RemoteSpace* space = static_cast<RemoteSpace*>(ptr->dl_tensor.data);
......@@ -181,6 +184,25 @@ class RPCModuleNode final : public ModuleNode {
PackedFunc fwrap_;
};
void* RPCWrappedFunc::UnwrapRemote(int rpc_sess_table_index,
const TVMArgValue& arg) {
if (arg.type_code() == kModuleHandle) {
Module mod = arg;
std::string tkey = mod->type_key();
CHECK_EQ(tkey, "rpc")
<< "ValueError: Cannot pass a non-RPC module to remote";
auto* rmod = static_cast<RPCModuleNode*>(mod.operator->());
CHECK_EQ(rmod->sess()->table_index(), rpc_sess_table_index)
<< "ValueError: Cannot pass in module into a different remote session";
return rmod->module_handle();
} else {
LOG(FATAL) << "ValueError: Cannot pass type "
<< runtime::TypeCode2Str(arg.type_code())
<< " as an argument to the remote";
return nullptr;
}
}
void RPCWrappedFunc::WrapRemote(std::shared_ptr<RPCSession> sess,
TVMArgs args,
TVMRetValue *rv) {
......
src/runtime/rpc/rpc_session.cc
......@@ -202,23 +202,33 @@ class RPCSession::EventHandler : public dmlc::Stream {
return ctx;
}
// Send Packed sequence to writer.
//
// client_mode: whether we are in client mode.
//
// funwrap: auxiliary function to unwrap remote Object
// when it is provided, we need to unwrap objects.
//
// return_ndarray is a special flag to handle returning of ndarray
// In this case, we return the shape, context and data of the array,
// as well as a customized PackedFunc that handles deletion of
// the array in the remote.
void SendPackedSeq(const TVMValue* arg_values,
const int* type_codes,
int n,
int num_args,
bool client_mode,
FUnwrapRemoteObject funwrap = nullptr,
bool return_ndarray = false) {
this->Write(n);
for (int i = 0; i < n; ++i) {
std::swap(client_mode_, client_mode);
this->Write(num_args);
for (int i = 0; i < num_args; ++i) {
int tcode = type_codes[i];
if (tcode == kNDArrayContainer) tcode = kArrayHandle;
this->Write(tcode);
}
// Argument packing.
for (int i = 0; i < n; ++i) {
for (int i = 0; i < num_args; ++i) {
int tcode = type_codes[i];
TVMValue value = arg_values[i];
switch (tcode) {
......@@ -241,7 +251,23 @@ class RPCSession::EventHandler : public dmlc::Stream {
break;
}
case kFuncHandle:
case kModuleHandle:
case kModuleHandle: {
// always send handle in 64 bit.
uint64_t handle;
// allow pass module as argument to remote.
if (funwrap != nullptr) {
void* remote_handle = (*funwrap)(
rpc_sess_table_index_,
runtime::TVMArgValue(value, tcode));
handle = reinterpret_cast<uint64_t>(remote_handle);
} else {
CHECK(!client_mode_)
<< "Cannot directly pass remote object as argument";
handle = reinterpret_cast<uint64_t>(value.v_handle);
}
this->Write(handle);
break;
}
case kHandle: {
// always send handle in 64 bit.
uint64_t handle = reinterpret_cast<uint64_t>(value.v_handle);
......@@ -300,6 +326,7 @@ class RPCSession::EventHandler : public dmlc::Stream {
}
}
}
std::swap(client_mode_, client_mode);
}
// Endian aware IO handling
......@@ -430,11 +457,11 @@ class RPCSession::EventHandler : public dmlc::Stream {
case kHandle:
case kStr:
case kBytes:
case kModuleHandle:
case kTVMContext: {
this->RequestBytes(sizeof(TVMValue)); break;
}
case kFuncHandle:
case kModuleHandle: {
case kFuncHandle: {
CHECK(client_mode_)
<< "Only client can receive remote functions";
this->RequestBytes(sizeof(TVMValue)); break;
......@@ -656,7 +683,7 @@ class RPCSession::EventHandler : public dmlc::Stream {
TVMValue ret_value;
ret_value.v_str = e.what();
int ret_tcode = kStr;
SendPackedSeq(&ret_value, &ret_tcode, 1);
SendPackedSeq(&ret_value, &ret_tcode, 1, false);
}
}
this->SwitchToState(kRecvCode);
......@@ -711,7 +738,7 @@ class RPCSession::EventHandler : public dmlc::Stream {
}
}
this->Write(code);
SendPackedSeq(&ret_value, &ret_tcode, 1);
SendPackedSeq(&ret_value, &ret_tcode, 1, false);
arg_recv_stage_ = 0;
this->SwitchToState(kRecvCode);
}
......@@ -734,7 +761,7 @@ class RPCSession::EventHandler : public dmlc::Stream {
if (rv.type_code() == kStr) {
ret_value.v_str = rv.ptr<std::string>()->c_str();
ret_tcode = kStr;
SendPackedSeq(&ret_value, &ret_tcode, 1);
SendPackedSeq(&ret_value, &ret_tcode, 1, false);
} else if (rv.type_code() == kBytes) {
std::string* bytes = rv.ptr<std::string>();
TVMByteArray arr;
......@@ -742,14 +769,14 @@ class RPCSession::EventHandler : public dmlc::Stream {
arr.size = bytes->length();
ret_value.v_handle = &arr;
ret_tcode = kBytes;
SendPackedSeq(&ret_value, &ret_tcode, 1);
SendPackedSeq(&ret_value, &ret_tcode, 1, false);
} else if (rv.type_code() == kFuncHandle ||
rv.type_code() == kModuleHandle) {
// always send handle in 64 bit.
CHECK(!client_mode_)
<< "Only server can send function and module handle back.";
rv.MoveToCHost(&ret_value, &ret_tcode);
SendPackedSeq(&ret_value, &ret_tcode, 1);
SendPackedSeq(&ret_value, &ret_tcode, 1, false);
} else if (rv.type_code() == kNDArrayContainer) {
// always send handle in 64 bit.
CHECK(!client_mode_)
......@@ -764,18 +791,18 @@ class RPCSession::EventHandler : public dmlc::Stream {
NDArray::Container* nd = static_cast<NDArray::Container*>(ret_value_pack[0].v_handle);
ret_value_pack[1].v_handle = nd;
ret_tcode_pack[1] = kHandle;
SendPackedSeq(ret_value_pack, ret_tcode_pack, 2, true);
SendPackedSeq(ret_value_pack, ret_tcode_pack, 2, false, nullptr, true);
} else {
ret_value = rv.value();
ret_tcode = rv.type_code();
SendPackedSeq(&ret_value, &ret_tcode, 1);
SendPackedSeq(&ret_value, &ret_tcode, 1, false);
}
} catch (const std::runtime_error& e) {
RPCCode code = RPCCode::kException;
this->Write(code);
ret_value.v_str = e.what();
ret_tcode = kStr;
SendPackedSeq(&ret_value, &ret_tcode, 1);
SendPackedSeq(&ret_value, &ret_tcode, 1, false);
}
}
......@@ -873,7 +900,7 @@ void RPCSession::Init() {
&reader_, &writer_, table_index_, name_, &remote_key_);
// Quick function to call remote.
call_remote_ = PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
handler_->SendPackedSeq(args.values, args.type_codes, args.num_args);
handler_->SendPackedSeq(args.values, args.type_codes, args.num_args, true);
RPCCode code = HandleUntilReturnEvent(rv, true, nullptr);
CHECK(code == RPCCode::kReturn) << "code=" << static_cast<int>(code);
});
......@@ -954,13 +981,16 @@ int RPCSession::ServerEventHandler(const std::string& bytes, int event_flag) {
void RPCSession::CallFunc(void* h,
TVMArgs args,
TVMRetValue* rv,
FUnwrapRemoteObject funwrap,
const PackedFunc* fwrap) {
std::lock_guard<std::recursive_mutex> lock(mutex_);
RPCCode code = RPCCode::kCallFunc;
handler_->Write(code);
uint64_t handle = reinterpret_cast<uint64_t>(h);
handler_->Write(handle);
handler_->SendPackedSeq(args.values, args.type_codes, args.num_args);
handler_->SendPackedSeq(
args.values, args.type_codes, args.num_args, true, funwrap);
code = HandleUntilReturnEvent(rv, true, fwrap);
CHECK(code == RPCCode::kReturn) << "code=" << static_cast<int>(code);
}
......
src/runtime/rpc/rpc_session.h
......@@ -91,6 +91,16 @@ enum class RPCCode : int {
};
/*!
* \brief Function that unwraps a remote object to its handle.
* \param rpc_sess_table_index RPC session table index for validation.
* \param obj Handle to the object argument.
* \return The corresponding handle.
*/
typedef void* (*FUnwrapRemoteObject)(
int rpc_sess_table_index,
const TVMArgValue& obj);
/*!
* \brief Abstract channel interface used to create RPCSession.
*/
class RPCChannel {
......@@ -144,11 +154,13 @@ class RPCSession {
* \param handle The function handle
* \param args The arguments
* \param rv The return value.
* \param funpwrap Function that takes a remote object and returns the raw handle.
* \param fwrap Wrapper function to turn Function/Module handle into real return.
*/
void CallFunc(RPCFuncHandle handle,
TVMArgs args,
TVMRetValue* rv,
FUnwrapRemoteObject funwrap,
const PackedFunc* fwrap);
/*!
* \brief Copy bytes into remote array content.
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment