tvm_tirx.runtime.Device

class tvm_tirx.runtime.Device

Bases: Device

TVM device strucure.

get_attr(name)

Query a named device attribute, returning None when unsupported.

exist returns a bool. max_thread_dimensions returns a JSON string; the property of the same name decodes it to a list. Vulkan exposes numeric version values as vulkan_api_version and vulkan_driver_version, separately from its formatted version strings.

property exist

Whether this device exists.

Returns True if TVM has support for the device, if the physical device is present, and the device is accessible through appropriate drivers (e.g. CUDA/Vulkan).

Returns:

exist (bool) – True if the device exists

property max_threads_per_block

Maximum number of threads on each block.

Returns device value for CUDA, Metal, ROCm, OpenCL, and Vulkan devices. Returns None for all other devices.

Returns:

max_threads_per_block (int or None) – The number of threads on each block

property warp_size

Number of threads that execute concurrently.

Returns device value for CUDA, ROCm, and Vulkan. Returns 1 for Metal and OpenCL devices, regardless of the physical device. Returns None for all other devices.

Returns:

warp_size (int or None) – Number of threads that execute concurrently

property max_shared_memory_per_block

Total amount of shared memory per block in bytes.

Returns device value for CUDA, ROCm, OpenCL, and Vulkan. Returns None for all other devices.

Returns:

max_shared_memory_per_block (int or None) – Total amount of shared memory per block in bytes

property compute_version

Get compute version number as string.

Returns maximum API version (e.g. CUDA/OpenCL/Vulkan) supported by the device.

Returns device value for CUDA, ROCm, OpenCL, and Vulkan. Returns None for all other devices.

Returns:

version (str or None) – The version string in major.minor format.

property device_name

Return the vendor-specific name of device.

Returns device value for CUDA, ROCm, OpenCL, and Vulkan. Returns None for all other devices.

Returns:

device_name (str or None) – The name of the device.

property max_clock_rate

Return the max clock frequency of device (kHz).

Returns device value for CUDA, ROCm, and OpenCL. Returns None for all other devices.

Returns:

max_clock_rate (int or None) – The maximum clock frequency of the device (kHz)

property multi_processor_count

Return the number of compute units in the device.

Returns device value for CUDA, ROCm, and OpenCL. Returns None for all other devices.

Returns:

multi_processor_count (int or None) – Thee number of compute units in the device

property max_thread_dimensions

Return the maximum size of each thread axis

Returns device value for CUDA, ROCm, OpenCL, and Vulkan. Returns None for all other devices.

Returns:

dims (List of int, or None) – The maximum length of threadIdx.x, threadIdx.y, threadIdx.z

property api_version

Returns version number of the SDK used to compile TVM.

For example, CUDA_VERSION for CUDA or VK_HEADER_VERSION for Vulkan.

Returns device value for CUDA, ROCm, OpenCL, and Vulkan. Returns None for all other devices.

Returns:

version (int or None) – The version of the SDK

property driver_version

Returns version number of the driver

Returns driver vendor’s internal version number. (e.g. “450.408.256” for nvidia-driver-450)

Returns device value for opencl and vulkan. Returns None for all other devices.

Returns:

version (str or None) – The version string in major.minor.patch format.

property l2_cache_size_bytes

Return the size of the device L2 cache in bytes

Supported devices include CUDA/ROCM/OpenCL.

Returns:

l2_cache_size_bytes (int or None) – The size of the device L2 cache in bytes returned by device runtime API. Return None if the device does not support this feature.

Note

The value returned by opencl’s API is smaller than actual device L2 cache size.

property total_global_memory

Return size of the total global memory.

Supported devices include CUDA/ROCm/Metal/OpenCL.

Returns:

total_global_memory (int or None) – Return the total size of global memory on device in bytes. Return None if the device does not support this feature.

property max_registers_per_block

Return the block register limit, or None when unsupported.

property gcn_arch

Return the ROCm architecture name, or None when unsupported.

property available_global_memory

Return size of the available global memory.

Supported devices include CUDA.

Returns:

available_global_memory (int or None) – Return the amount of unallocated global memory on device in bytes. Return None if the device does not support this feature.

create_raw_stream()

Create a new runtime stream at the context.

User should free the stream after use.

Returns:

stream (TVMStreamHandle) – The created runtime stream.

free_raw_stream(stream)

Free a created stream handle.

Parameters:

stream (TVMStreamHandle) – The stream which should to be released.

set_raw_stream(stream)

Set a created stream handle.

Parameters:

stream (TVMStreamHandle) – The stream which should to be set to the device.

sync(stream=None)

Synchronize until jobs finished at the context.

Parameters:

stream (TVMStreamHandle) – Jobs in this stream should be finished.