Tianqi Chen
acd48e9ade
* [RUNTIME] Enable ext_dev type for quick plugin of device * [TEST] Update testcase to cover all computation
169 lines
3.3 KiB
Python
169 lines
3.3 KiB
Python
"""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 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)
|