"""TVM Runtime NDArray API. tvm.ndarray provides a minimum runtime array API to test the correctness of the program. """ # pylint: disable=invalid-name,unused-import from __future__ import absolute_import as _abs import numpy as _np from ._ffi.ndarray import TVMContext, TVMType, NDArrayBase from ._ffi.ndarray import context, empty from ._ffi.ndarray import _set_class_ndarray from ._ffi.ndarray import register_extension, free_extension_handle class NDArray(NDArrayBase): """Lightweight NDArray class of TVM runtime. Strictly this is only an Array Container(a buffer object) No arthimetic operations are defined. All operations are performed by TVM functions. The goal is not to re-build yet another array library. Instead, this is a minimal data structure to demonstrate how can we use TVM in existing project which might have their own array containers. """ pass def cpu(dev_id=0): """Construct a CPU device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(1, dev_id) def gpu(dev_id=0): """Construct a CPU device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(2, dev_id) def rocm(dev_id=0): """Construct a ROCM device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(10, dev_id) def opencl(dev_id=0): """Construct a OpenCL device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(4, dev_id) def metal(dev_id=0): """Construct a metal device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(8, dev_id) def vpi(dev_id=0): """Construct a VPI simulated device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(9, dev_id) def opengl(dev_id=0): """Construct a OpenGL device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(11, dev_id) def ext_dev(dev_id=0): """Construct a extension device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context Note ---- This API is reserved for quick testing of new device by plugin device API as ext_dev. """ return TVMContext(12, dev_id) cl = opencl mtl = metal def array(arr, ctx=cpu(0)): """Create an array from source arr. Parameters ---------- arr : numpy.ndarray The array to be copied from ctx : TVMContext, optional The device context to create the array Returns ------- ret : NDArray The created array """ if not isinstance(arr, (_np.ndarray, NDArray)): arr = _np.array(arr) return empty(arr.shape, arr.dtype, ctx).copyfrom(arr) _set_class_ndarray(NDArray)