| """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_dlpack |
| 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 vulkan(dev_id=0): |
| """Construct a Vulkan device |
| |
| Parameters |
| ---------- |
| dev_id : int, optional |
| The integer device id |
| |
| Returns |
| ------- |
| ctx : TVMContext |
| The created context |
| """ |
| return TVMContext(7, 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) |
