[PYTHON][FFI] Cythonize NDArray.copyto

Author: tqchen

python/tvm/_ffi/_ctypes/ndarray.py

@@ -85,6 +85,12 @@ def __del__(self):
     def _tvm_handle(self):
         return ctypes.cast(self.handle, ctypes.c_void_p).value
 
+    def _copyto(self, target_nd):
+        """Internal function that implements copy to target ndarray."""
+        check_call(_LIB.TVMArrayCopyFromTo(
+                self.handle, target_nd.handle, None))
+        return target_nd
+
     def to_dlpack(self):
         """Produce an array from a DLPack Tensor without copying memory
 

python/tvm/_ffi/_cython/ndarray.pxi

@@ -76,6 +76,11 @@ cdef class NDArrayBase:
         if self.c_is_view == 0:
             CALL(TVMArrayFree(self.chandle))
 
+    def _copyto(self, target_nd):
+        """Internal function that implements copy to target ndarray."""
+        CALL(TVMArrayCopyFromTo(self.chandle, (<NDArrayBase>target_nd).chandle, NULL))
+        return target_nd
+
     def to_dlpack(self):
         """Produce an array from a DLPack Tensor without copying memory
 

python/tvm/_ffi/ndarray.py

@@ -294,14 +294,12 @@ def copyto(self, target):
         target : NDArray
             The target array to be copied, must have same shape as this array.
         """
-        if isinstance(target, TVMContext):
-            target = empty(self.shape, self.dtype, target)
         if isinstance(target, NDArrayBase):
-            check_call(_LIB.TVMArrayCopyFromTo(
-                self.handle, target.handle, None))
-        else:
-            raise ValueError("Unsupported target type %s" % str(type(target)))
-        return target
+            return self._copyto(target)
+        elif isinstance(target, TVMContext):
+            res = empty(self.shape, self.dtype, target)
+            return self._copyto(res)
+        raise ValueError("Unsupported target type %s" % str(type(target)))
 
 
 def free_extension_handle(handle, type_code):