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)