def webgpu(dev_id=0):
"""Construct a webgpu device.
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(15, dev_id)
def cuda(dev_id=0):
"""Construct a CUDA GPU device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(2, dev_id)
def vulkan(dev_id=0):
"""Construct a Vulkan device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(7, dev_id)
def hexagon(dev_id=0):
"""Construct a Hexagon device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(14, dev_id)
def vpi(dev_id=0):
"""Construct a VPI simulated device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(9, dev_id)
def metal(dev_id=0):
"""Construct a metal device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(8, dev_id)
def opencl(dev_id=0):
"""Construct a OpenCL device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(4, dev_id)
def rocm(dev_id=0):
"""Construct a ROCM device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(10, dev_id)
def micro_dev(dev_id=0):
"""Construct a micro device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
return Device(13, dev_id)
def ext_dev(dev_id=0):
"""Construct a extension device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
Note
----
This API is reserved for quick testing of new
device by plugin device API as ext_dev.
"""
return Device(12, dev_id)
def gpu(dev_id=0):
"""Construct a CUDA GPU device
deprecated:: 0.9.0
Use :py:func:`tvm.cuda` instead.
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
dev : Device
The created device
"""
warnings.warn(
"Please use tvm.cuda() instead of tvm.gpu(). tvm.gpu() is going to be deprecated in 0.9.0",
)
return Device(2, dev_id)
def device(dev_type, dev_id=0):
"""Construct a TVM device with given device type and id.
Parameters
----------
dev_type: int or str
The device type mask or name of the device.
dev_id : int, optional
The integer device id
Returns
-------
dev: tvm.runtime.Device
The corresponding device.
Examples
--------
Device can be used to create reflection of device by
string representation of the device type.
.. code-block:: python
assert tvm.device("cpu", 1) == tvm.cpu(1)
assert tvm.device("gpu", 0) == tvm.gpu(0)
assert tvm.device("cuda", 0) == tvm.gpu(0)
"""
if isinstance(dev_type, string_types):
if "-device=micro_dev" in dev_type:
dev_type = Device.STR2MASK["micro_dev"]
else:
dev_type = dev_type.split()[0]
if dev_type not in Device.STR2MASK:
raise ValueError("Unknown device type %s" % dev_type)
dev_type = Device.STR2MASK[dev_type]
return Device(dev_type, dev_id)