From 70a5902cd70cc451930454506bfb9cb1bb4e1f89 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Tue, 5 May 2020 12:07:59 -0700 Subject: [PATCH] [RPC] Call sync in remote cpu to gpu copies (#5512) --- src/runtime/rpc/minrpc/minrpc_server.h | 10 ++++++++++ src/runtime/rpc/rpc_local_session.cc | 6 ++++++ 2 files changed, 16 insertions(+) diff --git a/src/runtime/rpc/minrpc/minrpc_server.h b/src/runtime/rpc/minrpc/minrpc_server.h index 63ad359..a84042e 100644 --- a/src/runtime/rpc/minrpc/minrpc_server.h +++ b/src/runtime/rpc/minrpc/minrpc_server.h @@ -184,6 +184,11 @@ class MinRPCServer { data_ptr, 0, num_bytes, ctx, DLContext{kDLCPU, 0}, type_hint, nullptr); + // need sync to make sure that the copy is completed. + if (call_ecode == 0) { + call_ecode = TVMSynchronize( + ctx.device_type, ctx.device_id, nullptr); + } } if (call_ecode == 0) { @@ -223,6 +228,11 @@ class MinRPCServer { num_bytes, DLContext{kDLCPU, 0}, ctx, type_hint, nullptr); + // need sync to make sure that the copy is completed. + if (call_ecode == 0) { + call_ecode = TVMSynchronize( + ctx.device_type, ctx.device_id, nullptr); + } } if (call_ecode == 0) { diff --git a/src/runtime/rpc/rpc_local_session.cc b/src/runtime/rpc/rpc_local_session.cc index 0a2809b..351a989 100644 --- a/src/runtime/rpc/rpc_local_session.cc +++ b/src/runtime/rpc/rpc_local_session.cc @@ -98,6 +98,9 @@ void LocalSession::CopyToRemote(void* from, from, from_offset, to, to_offset, nbytes, cpu_ctx, ctx_to, type_hint, 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, @@ -115,6 +118,9 @@ void LocalSession::CopyFromRemote(void* from, from, from_offset, to, to_offset, nbytes, ctx_from, cpu_ctx, type_hint, nullptr); + // Copy can happen asynchrously + // synchronize to make sure that copy is completed + this->GetDeviceAPI(ctx_from)->StreamSync(ctx_from, nullptr); } void LocalSession::FreeHandle(void* handle, int type_code) { -- 2.7.4