tvm.ndarray

TVM Runtime NDArray API.

tvm.ndarray provides a minimum runtime array API to test the correctness of the program.

class tvm.ndarray.TVMContext(device_type, device_id)

TVM context strucure.

compute_version

Get compute verison number in string.

Currently used to get compute capability of CUDA device.

Returns:version – The version string in major.minor format.
Return type:str
device_name

Return the string name of device.

exist

Whether this device exist.

max_clock_rate

Return the max clock frequency of device.

max_shared_memory_per_block

Total amount of shared memory per block in bytes.

max_thread_dimensions

Return the maximum size of each thread axis

Returns:dims – The maximum length of threadIdx.x, threadIdx.y, threadIdx.z
Return type:List of int
max_threads_per_block

Maximum number of threads on each block.

multi_processor_count

Return the number of compute units of device.

sync()

Synchronize until jobs finished at the context.

warp_size

Number of threads that executes in concurrent.

class tvm.ndarray.NDArray(handle, is_view=False)

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.

asnumpy()

Convert this array to numpy array

Returns:np_arr – The corresponding numpy array.
Return type:numpy.ndarray
context

context of this array

copyfrom(source_array)

Peform an synchronize copy from the array.

Parameters:source_array (array_like) – The data source we should like to copy from.
Returns:arr – Reference to self.
Return type:NDArray
copyto(target)

Copy array to target

Parameters:target (NDArray) – The target array to be copied, must have same shape as this array.
ctx

context of this array

dtype

Type of this array

same_as(other)

Check object identity equality

Parameters:other (object) – The other object to compare to
Returns:same – Whether other is same as self.
Return type:bool
shape

Shape of this array

to_dlpack()

Produce an array from a DLPack Tensor without copying memory

Returns:dlpack
Return type:DLPack tensor view of the array data
tvm.cpu(dev_id=0)

Construct a CPU device

Parameters:dev_id (int, optional) – The integer device id
Returns:ctx – The created context
Return type:TVMContext
tvm.gpu(dev_id=0)

Construct a CPU device

Parameters:dev_id (int, optional) – The integer device id
Returns:ctx – The created context
Return type:TVMContext
tvm.opencl(dev_id=0)

Construct a OpenCL device

Parameters:dev_id (int, optional) – The integer device id
Returns:ctx – The created context
Return type:TVMContext
tvm.metal(dev_id=0)

Construct a metal device

Parameters:dev_id (int, optional) – The integer device id
Returns:ctx – The created context
Return type:TVMContext
tvm.ndarray.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 – The created array

Return type:

NDArray

tvm.ndarray.empty(shape, dtype='float32', ctx=cpu(0))

Create an empty array given shape and device

Parameters:
  • shape (tuple of int) – The shape of the array
  • dtype (type or str) – The data type of the array.
  • ctx (TVMContext) – The context of the array
Returns:

arr – The array tvm supported.

Return type:

tvm.nd.NDArray

tvm.register_extension(cls, fcreate=None)

Register a extension class to TVM.

After the class is registered, the class will be able to directly pass as Function argument generated by TVM.

Parameters:cls (class) – The class object to be registered as extension.

Note

The registered class is requires one property: _tvm_handle and a class attribute _tvm_tcode.

  • `_tvm_handle` returns integer represents the address of the handle.
  • `_tvm_tcode` gives integer represents type code of the class.
Returns:
  • cls (class) – The class being registered.
  • fcreate (function, optional) – The creation function to create a class object given handle value.

Example

The following code registers user defined class MyTensor to be DLTensor compatible.

@tvm.register_extension
class MyTensor(object):
    _tvm_tcode = tvm.TypeCode.ARRAY_HANDLE

    def __init__(self):
        self.handle = _LIB.NewDLTensor()

    @property
    def _tvm_handle(self):
        return self.handle.value