From 68ea2c3ed67fc4e007550703eedf5a2dcea5cff8 Mon Sep 17 00:00:00 2001
From: Chris Nuernberger <cnuernber@gmail.com>
Date: Sun, 4 Mar 2018 20:45:06 -0700
Subject: [PATCH] Small refactor for clarity in arraycopyfromto (#960)
---
src/runtime/c_runtime_api.cc | 27 +++++++++++++++------------
1 file changed, 15 insertions(+), 12 deletions(-)
diff --git a/src/runtime/c_runtime_api.cc b/src/runtime/c_runtime_api.cc
index 7d65a04bc..c7fd84362 100644
--- a/src/runtime/c_runtime_api.cc
+++ b/src/runtime/c_runtime_api.cc
@@ -413,19 +413,22 @@ int TVMArrayCopyFromTo(TVMArrayHandle from,
size_t from_size = GetDataSize(from);
size_t to_size = GetDataSize(to);
CHECK_EQ(from_size, to_size)
- << "TVMArrayCopyFromTo: The size must exactly match";
- TVMContext ctx = from->ctx;
- if (ctx.device_type == kDLCPU) {
- ctx = to->ctx;
- } else {
- CHECK(to->ctx.device_type == kDLCPU ||
- to->ctx.device_type == from->ctx.device_type)
- << "Can not copy across different ctx types directly";
- }
+ << "TVMArrayCopyFromTo: The size must exactly match";
+
+ CHECK(from->ctx.device_type == to->ctx.device_type
+ || from->ctx.device_type == kDLCPU
+ || to->ctx.device_type == kDLCPU)
+ << "Can not copy across different ctx types directly";
+
+ // Use the context that is *not* a cpu context to get the correct device
+ // api manager.
+ TVMContext ctx = from->ctx.device_type != kDLCPU ? from->ctx : to->ctx;
+
DeviceAPIManager::Get(ctx)->CopyDataFromTo(
- from->data, static_cast<size_t>(from->byte_offset),
- to->data, static_cast<size_t>(to->byte_offset),
- from_size, from->ctx, to->ctx, stream);
+ from->data, static_cast<size_t>(from->byte_offset),
+ to->data, static_cast<size_t>(to->byte_offset),
+ from_size, from->ctx, to->ctx, stream);
+
API_END();
}
--
GitLab