# 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)