From cbb0796e72b7ea70e9c7874ffaad0c0f966f8fae Mon Sep 17 00:00:00 2001
From: tqchen <tianqi.tchen@gmail.com>
Date: Wed, 13 Mar 2024 18:15:46 -0400
Subject: [PATCH] [RUNTIME] Ensure NDArray.CopyTo(Device) always sync

This PR ensures that NDArray.CopyTo(Device) always sync.
Prior to this PR, the behavior is uncertain as the underlying
DeviceAPI may or maynot sync. This PR further clarifies in
docs about the contract (that low-level device api is always async)
as well as the sync/async nature of each NDArray API.
---
 include/tvm/runtime/device_api.h |  2 ++
 include/tvm/runtime/ndarray.h    | 12 ++----------
 src/runtime/ndarray.cc           | 11 +++++++++++
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h
index 721990c625fa..b419212602c4 100644
--- a/include/tvm/runtime/device_api.h
+++ b/include/tvm/runtime/device_api.h
@@ -147,6 +147,8 @@ class TVM_DLL DeviceAPI {
    * \param from The source array.
    * \param to The target array.
    * \param stream Optional stream object.
+   * \note The copy may happen asynchronously if it involves a GPU context.
+   *       Call StreamSync to ensure the copy completes from host's pov.
    */
   virtual void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream);
   /*!
diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h
index 8400344bf559..d643355d2660 100644
--- a/include/tvm/runtime/ndarray.h
+++ b/include/tvm/runtime/ndarray.h
@@ -112,8 +112,9 @@ class NDArray : public ObjectRef {
    * \param dev The target device.
    * \param mem_scope The memory scope of the target array.
    * \return The array under another device.
+   * \note The copy always triggers a TVMSynchronize.
    */
-  inline NDArray CopyTo(const Device& dev, Optional<String> mem_scope = NullOpt) const;
+  TVM_DLL NDArray CopyTo(const Device& dev, Optional<String> mem_scope = NullOpt) const;
   /*!
    * \brief Load NDArray from stream
    * \param stream The input data stream
@@ -399,15 +400,6 @@ inline void NDArray::CopyTo(const NDArray& other) const {
   CopyFromTo(&(get_mutable()->dl_tensor), &(other.get_mutable()->dl_tensor));
 }
 
-inline NDArray NDArray::CopyTo(const Device& dev, Optional<String> mem_scope) const {
-  ICHECK(data_ != nullptr);
-  const DLTensor* dptr = operator->();
-  NDArray ret =
-      Empty(ShapeTuple(dptr->shape, dptr->shape + dptr->ndim), dptr->dtype, dev, mem_scope);
-  this->CopyTo(ret);
-  return ret;
-}
-
 inline int NDArray::use_count() const { return data_.use_count(); }
 
 inline const DLTensor* NDArray::operator->() const { return &(get_mutable()->dl_tensor); }
diff --git a/src/runtime/ndarray.cc b/src/runtime/ndarray.cc
index 675ee62a0511..6d03e2e01b51 100644
--- a/src/runtime/ndarray.cc
+++ b/src/runtime/ndarray.cc
@@ -287,6 +287,17 @@ void NDArray::CopyFromBytes(const void* data, size_t nbytes) {
   ArrayCopyFromBytes(&get_mutable()->dl_tensor, data, nbytes);
 }
 
+NDArray NDArray::CopyTo(const Device& dev, Optional<String> mem_scope) const {
+  ICHECK(data_ != nullptr);
+  const DLTensor* dptr = operator->();
+  NDArray ret =
+      Empty(ShapeTuple(dptr->shape, dptr->shape + dptr->ndim), dptr->dtype, dev, mem_scope);
+  this->CopyTo(ret);
+  Device copy_gpu_dev = dptr->device.device_type != kDLCPU ? dptr->device : dev;
+  DeviceAPI::Get(copy_gpu_dev)->StreamSync(copy_gpu_dev, nullptr);
+  return ret;
+}
+
 void NDArray::CopyFromTo(const DLTensor* from, DLTensor* to, TVMStreamHandle stream) {
   size_t from_size = GetDataSize(*from);
   size_t to_size = GetDataSize(*to);