def hexagon(dev_id=0): """Construct a Hexagon device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(14, dev_id)
def vulkan(dev_id=0): """Construct a Vulkan device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(7, dev_id)
def micro_dev(dev_id=0): """Construct a micro device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(13, dev_id)
def opengl(dev_id=0): """Construct a OpenGL device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(11, dev_id)
def vpi(dev_id=0): """Construct a VPI simulated device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(9, dev_id)
def rocm(dev_id=0): """Construct a ROCM device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(10, dev_id)
def gpu(dev_id=0): """Construct a GPU device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(2, dev_id)
def webgpu(dev_id=0): """Construct a webgpu device. Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context """ return TVMContext(15, dev_id)
def ext_dev(dev_id=0): """Construct a extension device Parameters ---------- dev_id : int, optional The integer device id Returns ------- ctx : TVMContext The created context Note ---- This API is reserved for quick testing of new device by plugin device API as ext_dev. """ return TVMContext(12, dev_id)
def context(dev_type, dev_id=0): """Construct a TVM context 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 ------- ctx: tvm.runtime.TVMContext The corresponding context. Examples -------- Context can be used to create reflection of context by string representation of the device type. .. code-block:: python assert tvm.context("cpu", 1) == tvm.cpu(1) assert tvm.context("gpu", 0) == tvm.gpu(0) assert tvm.context("cuda", 0) == tvm.gpu(0) """ if isinstance(dev_type, string_types): if '-device=micro_dev' in dev_type: dev_type = 'micro_dev' else: dev_type = dev_type.split()[0] if dev_type not in TVMContext.STR2MASK: raise ValueError("Unknown device type %s" % dev_type) dev_type = TVMContext.STR2MASK[dev_type] return TVMContext(dev_type, dev_id)