tqchen commented on a change in pull request #7488:
URL: https://github.com/apache/tvm/pull/7488#discussion_r581040168
##########
File path: src/runtime/rpc/rpc_device_api.cc
##########
@@ -65,30 +77,41 @@ class RPCDeviceAPI final : public DeviceAPI {
}
delete space;
}
- void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t
to_offset, size_t size,
- TVMContext ctx_from, TVMContext ctx_to, DLDataType
type_hint,
- TVMStreamHandle stream) final {
+
+ void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream)
final {
+ DLContext ctx_from = from->ctx;
+ DLContext ctx_to = to->ctx;
if (IsRPCSessionContext(ctx_from) && IsRPCSessionContext(ctx_to)) {
ICHECK(ctx_from.device_type == ctx_to.device_type)
<< "Cannot copy across two different remote session";
auto remote_ctx_from = RemoveRPCSessionMask(ctx_from);
auto remote_ctx_to = RemoveRPCSessionMask(ctx_to);
+ void* from_data = static_cast<const RemoteSpace*>(from->data)->data;
+ void* to_data = static_cast<const RemoteSpace*>(to->data)->data;
+ DLTensor from_tensor{from_data, remote_ctx_from, from->ndim,
from->dtype,
Review comment:
DLTensor from_tensor = *from;
from_tensor.ctx = RemoveRPCSessionMask(ctx_from);
from_tensor.data = static_cast<const RemoteSpace*>(from->data)->data;
##########
File path: src/runtime/rpc/rpc_endpoint.cc
##########
@@ -904,32 +883,41 @@ void RPCDevAllocData(RPCSession* handler, TVMArgs args,
TVMRetValue* rv) {
*rv = data;
}
+void RPCDevAllocDataWithScope(RPCSession* handler, TVMArgs args, TVMRetValue*
rv) {
+ DLTensor* arr = args[0];
+ TVMContext ctx = arr->ctx;
+ std::vector<int64_t> shape(arr->shape, arr->shape + arr->ndim);
+ DLDataType dtype = arr->dtype;
+ int tcode = args[1].type_code();
+ Optional<String> mem_scope = NullOpt;
+ if (tcode == kTVMStr) {
+ mem_scope = args[1].operator std::string();
Review comment:
consider `operator String()`(although std::string also works)
##########
File path: src/runtime/rpc/rpc_endpoint.cc
##########
@@ -815,51 +798,47 @@ void RPCEndpoint::CallFunc(RPCSession::PackedFuncHandle
h, const TVMValue* arg_v
ICHECK(code == RPCCode::kReturn) << "code=" << static_cast<int>(code);
}
-void RPCEndpoint::CopyToRemote(void* from, size_t from_offset, void* to,
size_t to_offset,
- size_t data_size, TVMContext ctx_to, DLDataType
type_hint) {
+void RPCEndpoint::CopyToRemote(void* from_bytes, DLTensor* to, uint64_t
nbytes) {
std::lock_guard<std::mutex> lock(mutex_);
RPCCode code = RPCCode::kCopyToRemote;
- uint64_t handle = reinterpret_cast<uint64_t>(to);
- uint64_t offset = static_cast<uint64_t>(to_offset);
- uint64_t size = static_cast<uint64_t>(data_size);
- uint64_t packet_nbytes = sizeof(code) + sizeof(handle) + sizeof(offset) +
sizeof(size) +
- sizeof(ctx_to) + sizeof(type_hint) + data_size;
+ uint64_t num_data_bytes = static_cast<uint64_t>(GetDataSize(*to));
+ ICHECK_EQ(nbytes, num_data_bytes);
+
+ uint64_t to_data = reinterpret_cast<uint64_t>(to->data);
+ uint64_t num_shape_bytes = to->ndim * sizeof(int64_t);
Review comment:
shape_bytes
##########
File path: src/runtime/rpc/rpc_endpoint.cc
##########
@@ -815,51 +798,47 @@ void RPCEndpoint::CallFunc(RPCSession::PackedFuncHandle
h, const TVMValue* arg_v
ICHECK(code == RPCCode::kReturn) << "code=" << static_cast<int>(code);
}
-void RPCEndpoint::CopyToRemote(void* from, size_t from_offset, void* to,
size_t to_offset,
- size_t data_size, TVMContext ctx_to, DLDataType
type_hint) {
+void RPCEndpoint::CopyToRemote(void* from_bytes, DLTensor* to, uint64_t
nbytes) {
std::lock_guard<std::mutex> lock(mutex_);
RPCCode code = RPCCode::kCopyToRemote;
- uint64_t handle = reinterpret_cast<uint64_t>(to);
- uint64_t offset = static_cast<uint64_t>(to_offset);
- uint64_t size = static_cast<uint64_t>(data_size);
- uint64_t packet_nbytes = sizeof(code) + sizeof(handle) + sizeof(offset) +
sizeof(size) +
- sizeof(ctx_to) + sizeof(type_hint) + data_size;
+ uint64_t num_data_bytes = static_cast<uint64_t>(GetDataSize(*to));
+ ICHECK_EQ(nbytes, num_data_bytes);
+
+ uint64_t to_data = reinterpret_cast<uint64_t>(to->data);
+ uint64_t num_shape_bytes = to->ndim * sizeof(int64_t);
+ uint64_t packet_nbytes = sizeof(code) + sizeof(to_data) + sizeof(to->ctx) +
sizeof(to->ndim) +
+ sizeof(to->dtype) + sizeof(to->byte_offset) +
num_shape_bytes +
+ sizeof(nbytes) + num_data_bytes;
handler_->Write(packet_nbytes);
handler_->Write(code);
- handler_->Write(handle);
- handler_->Write(offset);
- handler_->Write(size);
- handler_->Write(ctx_to);
- handler_->Write(type_hint);
- handler_->WriteArray(reinterpret_cast<char*>(from) + from_offset, data_size);
-
+ RPCReference::SendDLTensor(handler_, to);
+ handler_->Write(nbytes);
+ handler_->WriteArray(reinterpret_cast<char*>(from_bytes), nbytes);
ICHECK(HandleUntilReturnEvent(true, [](TVMArgs) {}) == RPCCode::kReturn);
}
-void RPCEndpoint::CopyFromRemote(void* from, size_t from_offset, void* to,
size_t to_offset,
- size_t data_size, TVMContext ctx_from,
DLDataType type_hint) {
+void RPCEndpoint::CopyFromRemote(DLTensor* from, void* to_bytes, uint64_t
nbytes) {
std::lock_guard<std::mutex> lock(mutex_);
RPCCode code = RPCCode::kCopyFromRemote;
- uint64_t handle = reinterpret_cast<uint64_t>(from);
- uint64_t offset = static_cast<uint64_t>(from_offset);
- uint64_t size = static_cast<uint64_t>(data_size);
- uint64_t packet_nbytes = sizeof(code) + sizeof(handle) + sizeof(offset) +
sizeof(size) +
- sizeof(ctx_from) + sizeof(type_hint);
+ uint64_t num_data_bytes = static_cast<uint64_t>(GetDataSize(*from));
+ CHECK_EQ(nbytes, num_data_bytes);
+
+ uint64_t from_data = reinterpret_cast<uint64_t>(from->data);
+ uint64_t num_shape_bytes = from->ndim * sizeof(int64_t);
Review comment:
shape_bytes
##########
File path: src/runtime/rpc/rpc_endpoint.h
##########
@@ -135,8 +135,7 @@ class RPCEndpoint {
* \param ctx_to The target context.
* \param type_hint Hint of content data type.
*/
- void CopyToRemote(void* from, size_t from_offset, void* to, size_t
to_offset, size_t nbytes,
- TVMContext ctx_to, DLDataType type_hint);
+ void CopyToRemote(void* from_bytes, DLTensor* to, uint64_t nbytes);
Review comment:
consistency: either use size_t nbytes or uint64_t nbytes everywhere
(RPCSession and RPCEndpoint)
##########
File path: src/runtime/rpc/rpc_local_session.cc
##########
@@ -87,26 +88,25 @@ void LocalSession::CallFunc(RPCSession::PackedFuncHandle
func, const TVMValue* a
this->EncodeReturn(std::move(rv), encode_return);
}
-void LocalSession::CopyToRemote(void* from, size_t from_offset, void* to,
size_t to_offset,
- size_t nbytes, TVMContext ctx_to, DLDataType
type_hint) {
- TVMContext cpu_ctx;
- cpu_ctx.device_type = kDLCPU;
- cpu_ctx.device_id = 0;
- this->GetDeviceAPI(ctx_to)->CopyDataFromTo(from, from_offset, to, to_offset,
nbytes, cpu_ctx,
- ctx_to, type_hint, nullptr);
+void LocalSession::CopyToRemote(void* from_bytes, DLTensor* to, size_t nbytes)
{
+ ICHECK_EQ(nbytes, GetDataSize(*to));
+ std::vector<int64_t> shape(to->shape, to->shape + to->ndim);
+ DLTensor from{from_bytes, {kDLCPU, 0}, to->ndim, to->dtype, shape.data(),
nullptr, 0};
+
+ TVMContext ctx_to = to->ctx;
+ this->GetDeviceAPI(ctx_to)->CopyDataFromTo(&from, to, nullptr);
// Copy can happen asynchrously
// synchronize to make sure that copy is completed
this->GetDeviceAPI(ctx_to)->StreamSync(ctx_to, nullptr);
}
-void LocalSession::CopyFromRemote(void* from, size_t from_offset, void* to,
size_t to_offset,
- size_t nbytes, TVMContext ctx_from,
DLDataType type_hint) {
- TVMContext cpu_ctx;
- cpu_ctx.device_type = kDLCPU;
- cpu_ctx.device_id = 0;
+void LocalSession::CopyFromRemote(DLTensor* from, void* to_bytes, size_t
nbytes) {
+ ICHECK_EQ(nbytes, GetDataSize(*from));
+ std::vector<int64_t> shape(from->shape, from->shape + from->ndim);
Review comment:
no need to construct a vector, can directlyr euse from->shape and
from->ndim
##########
File path: src/runtime/rpc/rpc_local_session.cc
##########
@@ -87,26 +88,25 @@ void LocalSession::CallFunc(RPCSession::PackedFuncHandle
func, const TVMValue* a
this->EncodeReturn(std::move(rv), encode_return);
}
-void LocalSession::CopyToRemote(void* from, size_t from_offset, void* to,
size_t to_offset,
- size_t nbytes, TVMContext ctx_to, DLDataType
type_hint) {
- TVMContext cpu_ctx;
- cpu_ctx.device_type = kDLCPU;
- cpu_ctx.device_id = 0;
- this->GetDeviceAPI(ctx_to)->CopyDataFromTo(from, from_offset, to, to_offset,
nbytes, cpu_ctx,
- ctx_to, type_hint, nullptr);
+void LocalSession::CopyToRemote(void* from_bytes, DLTensor* to, size_t nbytes)
{
+ ICHECK_EQ(nbytes, GetDataSize(*to));
+ std::vector<int64_t> shape(to->shape, to->shape + to->ndim);
+ DLTensor from{from_bytes, {kDLCPU, 0}, to->ndim, to->dtype, shape.data(),
nullptr, 0};
Review comment:
construct DLTensor by fields.
##########
File path: web/emcc/tvmjs_support.cc
##########
@@ -177,33 +177,25 @@ class AsyncLocalSession : public LocalSession {
}
}
- void AsyncCopyToRemote(void* local_from, size_t local_from_offset, void*
remote_to,
- size_t remote_to_offset, size_t nbytes, TVMContext
remote_ctx_to,
- DLDataType type_hint, FAsyncCallback on_complete)
final {
- TVMContext cpu_ctx;
- cpu_ctx.device_type = kDLCPU;
- cpu_ctx.device_id = 0;
+ void AsyncCopyToRemote(void* local_from_bytes, DLTensor* remote_to, size_t
nbytes,
+ FAsyncCallback on_complete) final {
try {
- this->GetDeviceAPI(remote_ctx_to)
- ->CopyDataFromTo(local_from, local_from_offset, remote_to,
remote_to_offset, nbytes,
- cpu_ctx, remote_ctx_to, type_hint, nullptr);
- this->AsyncStreamWait(remote_ctx_to, nullptr, on_complete);
+ this->GetDeviceAPI(remote_to->ctx)
Review comment:
update to the new API that uses DLTensor*
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
