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 mem_scope = NullOpt) const; + TVM_DLL NDArray CopyTo(const Device& dev, Optional 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 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 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);