wehub-resource-sync
331 lines
10 KiB
Python
331 lines
10 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Common runtime ctypes."""
|
|
|
|
# pylint: disable=invalid-name
|
|
import json
|
|
|
|
import tvm_ffi
|
|
|
|
from . import _ffi_api
|
|
|
|
RPC_SESS_MASK = 128
|
|
|
|
|
|
class Device(tvm_ffi.core.Device):
|
|
"""TVM device strucure."""
|
|
|
|
def _GetDeviceAttr(self, device_type, device_id, attr_id):
|
|
"""Internal helper function to invoke runtime.GetDeviceAttr"""
|
|
# pylint: disable=import-outside-toplevel
|
|
return _ffi_api.GetDeviceAttr(device_type, device_id, attr_id)
|
|
|
|
@property
|
|
def exist(self):
|
|
"""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
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 0) != 0
|
|
|
|
@property
|
|
def max_threads_per_block(self):
|
|
"""Maximum number of threads on each block.
|
|
|
|
Returns device value for CUDA, Metal, ROCm, OpenCL, and Vulkan
|
|
devices. Returns remote device value for RPC devices.
|
|
Returns None for all other devices.
|
|
|
|
Returns
|
|
-------
|
|
max_threads_per_block : int or None
|
|
The number of threads on each block
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 1)
|
|
|
|
@property
|
|
def warp_size(self):
|
|
"""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 remote device value for RPC devices. Returns
|
|
None for all other devices.
|
|
|
|
Returns
|
|
-------
|
|
warp_size : int or None
|
|
Number of threads that execute concurrently
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 2)
|
|
|
|
@property
|
|
def max_shared_memory_per_block(self):
|
|
"""Total amount of shared memory per block in bytes.
|
|
|
|
Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
|
|
Returns remote device value for RPC devices. Returns None for
|
|
all other devices.
|
|
|
|
Returns
|
|
-------
|
|
max_shared_memory_per_block : int or None
|
|
Total amount of shared memory per block in bytes
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 3)
|
|
|
|
@property
|
|
def compute_version(self):
|
|
"""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 remote device value for RPC devices. Returns
|
|
None for all other devices.
|
|
|
|
Returns
|
|
-------
|
|
version : str or None
|
|
The version string in `major.minor` format.
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 4)
|
|
|
|
@property
|
|
def device_name(self):
|
|
"""Return the vendor-specific name of device.
|
|
|
|
Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
|
|
Returns remote device value for RPC devices. Returns None for
|
|
all other devices.
|
|
|
|
Returns
|
|
-------
|
|
device_name : str or None
|
|
The name of the device.
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 5)
|
|
|
|
@property
|
|
def max_clock_rate(self):
|
|
"""Return the max clock frequency of device (kHz).
|
|
|
|
Returns device value for CUDA, ROCm, and OpenCL. Returns
|
|
remote device value for RPC devices. Returns None for all
|
|
other devices.
|
|
|
|
Returns
|
|
-------
|
|
max_clock_rate : int or None
|
|
The maximum clock frequency of the device (kHz)
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 6)
|
|
|
|
@property
|
|
def multi_processor_count(self):
|
|
"""Return the number of compute units in the device.
|
|
|
|
Returns device value for CUDA, ROCm, and OpenCL. Returns
|
|
remote device value for RPC devices. Returns None for all
|
|
other devices.
|
|
|
|
Returns
|
|
-------
|
|
multi_processor_count : int or None
|
|
Thee number of compute units in the device
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 7)
|
|
|
|
@property
|
|
def max_thread_dimensions(self):
|
|
"""Return the maximum size of each thread axis
|
|
|
|
Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
|
|
Returns remote device value for RPC devices. Returns None for
|
|
all other devices.
|
|
|
|
Returns
|
|
-------
|
|
dims: List of int, or None
|
|
The maximum length of threadIdx.x, threadIdx.y, threadIdx.z
|
|
|
|
"""
|
|
return json.loads(self._GetDeviceAttr(self.dlpack_device_type(), self.index, 8))
|
|
|
|
@property
|
|
def api_version(self):
|
|
"""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 remote device value for RPC devices. Returns None for
|
|
all other devices.
|
|
|
|
Returns
|
|
-------
|
|
version : int or None
|
|
The version of the SDK
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 11)
|
|
|
|
@property
|
|
def driver_version(self):
|
|
"""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 remote
|
|
device value for RPC devices. Returns None for all other
|
|
devices.
|
|
|
|
Returns
|
|
-------
|
|
version : str or None
|
|
The version string in `major.minor.patch` format.
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 12)
|
|
|
|
@property
|
|
def l2_cache_size_bytes(self):
|
|
"""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.
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 13)
|
|
|
|
@property
|
|
def total_global_memory(self):
|
|
"""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.
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 14)
|
|
|
|
@property
|
|
def available_global_memory(self):
|
|
"""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.
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 15)
|
|
|
|
def texture_spatial_limit(self):
|
|
"""Returns limits for textures by spatial dimensions
|
|
|
|
Returns
|
|
-------
|
|
limit : int or None
|
|
Maximum size of the texture by spatial dimensions
|
|
|
|
"""
|
|
return self._GetDeviceAttr(self.dlpack_device_type(), self.index, 12)
|
|
|
|
def create_raw_stream(self):
|
|
"""Create a new runtime stream at the context.
|
|
|
|
User should free the stream after use.
|
|
|
|
Returns
|
|
-------
|
|
stream : TVMStreamHandle
|
|
The created runtime stream.
|
|
"""
|
|
return _ffi_api.Device_StreamCreate(self)
|
|
|
|
def free_raw_stream(self, stream):
|
|
"""Free a created stream handle.
|
|
|
|
Parameters
|
|
----------
|
|
stream : TVMStreamHandle
|
|
The stream which should to be released.
|
|
"""
|
|
_ffi_api.Device_StreamFree(self, stream)
|
|
|
|
def set_raw_stream(self, stream):
|
|
"""Set a created stream handle.
|
|
|
|
Parameters
|
|
----------
|
|
stream : TVMStreamHandle
|
|
The stream which should to be set to the device.
|
|
"""
|
|
_ffi_api.Device_SetStream(self, stream)
|
|
|
|
def sync(self, stream=None):
|
|
"""Synchronize until jobs finished at the context.
|
|
|
|
Parameters
|
|
----------
|
|
stream : TVMStreamHandle
|
|
Jobs in this stream should be finished.
|
|
"""
|
|
_ffi_api.Device_StreamSync(self, stream or 0)
|
|
|
|
def __device_type_name__(self):
|
|
if self.dlpack_device_type() >= RPC_SESS_MASK:
|
|
tbl_id = self.dlpack_device_type() / RPC_SESS_MASK - 1
|
|
dev_type = self.dlpack_device_type() % RPC_SESS_MASK
|
|
return f"remote[{tbl_id}]:{Device._DEVICE_TYPE_TO_NAME[dev_type]}"
|
|
return Device._DEVICE_TYPE_TO_NAME[self.dlpack_device_type()]
|
|
|
|
|
|
tvm_ffi.core._set_class_device(Device)
|