def build_target_by_device_type_map(target):
"""Build a map from DLDevice device_type to a Target used with that device.
At runtime, TVM assigns target code to DLDevices by determining a device_type for each Target.
This function handles this process at compile time and, as a side effect, validates that exactly
one target maps to one device_type.
Parameters
----------
target : Target or str or dict
If a Target or str: assumes that exactly one device type is present in the model.
If a dict: keys are tvm.ndarray.device, values are the targets used for each device.
Returns
-------
"""
target = target if target else Target.current()
if target is None:
raise ValueError("Target is not set in env or passed as argument.")
tgts = {}
if isinstance(target, (str, Target)):
dev_type = tvm_expr.IntImm("int32",
_nd.device(str(target)).device_type)
tgts[dev_type] = Target(target)
elif isinstance(target, dict):
for dev, tgt in target.items():
dev_type = tvm_expr.IntImm("int32", _nd.device(dev).device_type)
tgts[dev_type] = Target(tgt)
else:
raise TypeError("target is expected to be str or " +
"tvm.target.Target, but received " +
"{}".format(type(target)))
return tgts
def _update_target(target):
target = target if target else _target.Target.current()
if target is None:
raise ValueError("Target is not set in env or passed as argument.")
tgts = {}
if isinstance(target, (str, _target.Target)):
dev_type = tvm_expr.IntImm("int32", _nd.context(str(target)).device_type)
tgts[dev_type] = _target.create(target)
elif isinstance(target, dict):
for dev, tgt in target.items():
dev_type = tvm_expr.IntImm("int32", _nd.context(dev).device_type)
tgts[dev_type] = _target.create(tgt)
else:
raise TypeError("target is expected to be str or " +
"tvm.target.Target, but received " +
"{}".format(type(target)))
return tgts
def _setup(self, mod, target):
tgts = {}
if isinstance(target, dict):
for dev, tgt in target.items():
if not isinstance(tgt, (str, Target)):
raise Exception("Unknown target type")
tgts[dev] = Target(tgt)
elif isinstance(target, (str, Target)):
tgts[_expr.IntImm("int32", 0)] = Target(target)
self._init(mod, tgts)
from tvm.tir import expr as _expr
a = _expr.IntImm("int32", 1)
print(type(a.dtype))
print(type(a.value))
class b():
def __init__(self, value):
self.value = value
res = b(123)
print(dir(b))
