| # 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. |
| from enum import IntEnum |
| from numbers import Integral |
| from typing import Any, Optional |
| |
| _CLASS_DEVICE = None |
| |
| |
| def _set_class_device(cls): |
| global _CLASS_DEVICE |
| _CLASS_DEVICE = cls |
| |
| |
| def _create_device_from_tuple(cls, device_type, device_id): |
| cdef DLDevice cdevice = TVMFFIDLDeviceFromIntPair(device_type, device_id) |
| ret = cls.__new__(cls) |
| (<Device>ret).cdevice = cdevice |
| return ret |
| |
| |
| class DLDeviceType(IntEnum): |
| """Enumeration mirroring DLPack's `DLDeviceType <https://dmlc.github.io/dlpack/latest/c_api.html#c.DLDeviceType>`_ |
| |
| Values can be compared against :py:meth:`Device.dlpack_device_type`. |
| |
| Examples |
| -------- |
| .. code-block:: python |
| |
| import tvm_ffi |
| dev = tvm_ffi.device("cuda", 0) |
| assert dev.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCUDA |
| |
| """ |
| kDLCPU = 1 |
| kDLCUDA = 2 |
| kDLCUDAHost = 3 |
| kDLOpenCL = 4 |
| kDLVulkan = 7 |
| kDLMetal = 8 |
| kDLVPI = 9 |
| kDLROCM = 10 |
| kDLROCMHost = 11 |
| kDLExtDev = 12 |
| kDLCUDAManaged = 13 |
| kDLOneAPI = 14 |
| kDLWebGPU = 15 |
| kDLHexagon = 16 |
| kDLMAIA = 17 |
| kDLTrn = 18 |
| |
| |
| cdef class Device: |
| """A device descriptor used by TVM FFI and DLPack. |
| |
| A :class:`Device` identifies a placement (e.g. CPU, CUDA GPU) and |
| a device index within that placement. Most users construct devices |
| using :func:`tvm_ffi.device`. |
| |
| Parameters |
| ---------- |
| device_type : Union[str, int] |
| A device type name (e.g. ``"cpu"``, ``"cuda"``) or a |
| DLPack device type code. |
| index : int, optional |
| Zero-based device index (defaults to ``0`` when omitted). |
| |
| Examples |
| -------- |
| .. code-block:: python |
| |
| import tvm_ffi |
| |
| dev = tvm_ffi.device("cuda:0") |
| assert dev.type == "cuda" |
| assert dev.index == 0 |
| assert str(dev) == "cuda:0" |
| |
| """ |
| __slots__ = () |
| cdef DLDevice cdevice |
| |
| _DEVICE_TYPE_TO_NAME = { |
| DLDeviceType.kDLCPU: "cpu", |
| DLDeviceType.kDLCUDA: "cuda", |
| DLDeviceType.kDLCUDAHost: "cuda_host", |
| DLDeviceType.kDLCUDAManaged: "cuda_managed", |
| DLDeviceType.kDLOpenCL: "opencl", |
| DLDeviceType.kDLVulkan: "vulkan", |
| DLDeviceType.kDLMetal: "metal", |
| DLDeviceType.kDLVPI: "vpi", |
| DLDeviceType.kDLROCM: "rocm", |
| DLDeviceType.kDLROCMHost: "rocm_host", |
| DLDeviceType.kDLExtDev: "ext_dev", |
| DLDeviceType.kDLOneAPI: "oneapi", |
| DLDeviceType.kDLWebGPU: "webgpu", |
| DLDeviceType.kDLHexagon: "hexagon", |
| DLDeviceType.kDLMAIA: "maia", |
| DLDeviceType.kDLTrn: "trn", |
| } |
| |
| _DEVICE_NAME_TO_TYPE = { |
| "llvm": DLDeviceType.kDLCPU, |
| "cpu": DLDeviceType.kDLCPU, |
| "c": DLDeviceType.kDLCPU, |
| "test": DLDeviceType.kDLCPU, |
| "cuda": DLDeviceType.kDLCUDA, |
| "nvptx": DLDeviceType.kDLCUDA, |
| "cl": DLDeviceType.kDLOpenCL, |
| "opencl": DLDeviceType.kDLOpenCL, |
| "vulkan": DLDeviceType.kDLVulkan, |
| "metal": DLDeviceType.kDLMetal, |
| "vpi": DLDeviceType.kDLVPI, |
| "rocm": DLDeviceType.kDLROCM, |
| "ext_dev": DLDeviceType.kDLExtDev, |
| "hexagon": DLDeviceType.kDLHexagon, |
| "webgpu": DLDeviceType.kDLWebGPU, |
| "maia": DLDeviceType.kDLMAIA, |
| "trn": DLDeviceType.kDLTrn, |
| } |
| |
| def __init__(self, device_type: str | int, index: Optional[Integral] = None) -> None: |
| device_type_or_name = device_type |
| index = index if index is not None else 0 |
| if isinstance(device_type_or_name, str): |
| # skip suffix annotations |
| device_type_or_name = device_type_or_name.split(" ")[0] |
| parts = device_type_or_name.split(":") |
| if len(parts) < 1 or len(parts) > 2: |
| raise ValueError(f"Invalid device: {device_type_or_name}") |
| if parts[0] not in self._DEVICE_NAME_TO_TYPE: |
| raise ValueError(f"Unknown device: {parts[0]}") |
| device_type = self._DEVICE_NAME_TO_TYPE[parts[0]] |
| if len(parts) == 2: |
| try: |
| index = int(parts[1]) |
| except ValueError: |
| raise ValueError(f"Invalid device index: {parts[1]}") |
| else: |
| device_type = device_type_or_name |
| |
| if not isinstance(index, Integral): |
| if hasattr(index, "item") and callable(index.item): |
| index = index.item() |
| if not isinstance(index, Integral): |
| raise TypeError(f"Invalid device index: {index}") |
| self.cdevice = TVMFFIDLDeviceFromIntPair(device_type, index) |
| |
| def __reduce__(self) -> Any: |
| cls = type(self) |
| return (_create_device_from_tuple, (cls, self.cdevice.device_type, self.cdevice.device_id)) |
| |
| def __eq__(self, other: object) -> bool: |
| if not isinstance(other, Device): |
| return False |
| return ( |
| self.cdevice.device_type == (<Device>other).cdevice.device_type |
| and self.cdevice.device_id == (<Device>other).cdevice.device_id |
| ) |
| |
| def __ne__(self, other: object) -> bool: |
| return not self.__eq__(other) |
| |
| def __str__(self) -> str: |
| cdef int dev_type = self.cdevice.device_type |
| name = self.__device_type_name__() |
| index = self.cdevice.device_id |
| return f"{name}:{index}" |
| |
| def __repr__(self) -> str: |
| cdef int dev_type = self.cdevice.device_type |
| name = self.__device_type_name__() |
| index = self.cdevice.device_id |
| return f"device(type='{name}', index={index})" |
| |
| def __hash__(self) -> int: |
| return hash((self.cdevice.device_type, self.cdevice.device_id)) |
| |
| def __device_type_name__(self) -> str: |
| """Return the canonical device type name (e.g. ``"cuda"``).""" |
| return self._DEVICE_TYPE_TO_NAME[self.cdevice.device_type] |
| |
| @property |
| def type(self) -> str: |
| """Device type name such as ``"cpu"`` or ``"cuda"``.""" |
| return self.__device_type_name__() |
| |
| @property |
| def index(self) -> int: |
| """Zero-based device index.""" |
| return self.cdevice.device_id |
| |
| def dlpack_device_type(self) -> int: |
| """Return the corresponding :class:`DLDeviceType` enum value.""" |
| return self.cdevice.device_type |
| |
| |
| cdef inline object make_ret_device(TVMFFIAny result): |
| ret = _CLASS_DEVICE.__new__(_CLASS_DEVICE) |
| (<Device>ret).cdevice = result.v_device |
| return ret |
| |
| |
| _set_class_device(Device) |
