def _update_target(target):
target = target if target else _target.current_target()
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.Target)):
raise Exception("Unknown target type")
tgts[dev] = _target.create(tgt)
elif isinstance(target, (str, _target.Target)):
tgts[_expr.IntImm("int32", 0)] = _target.create(target)
self._init(mod, tgts)
def after_C_priority():
input_1 = relay.var('input_1', shape=(10, 10))
input_2 = relay.var('input_2', shape=(10, 10))
add = relay.add(input_1, input_2)
x = relay.var('x')
out = relay.abs(x)
out = relay.nn.relu(out)
merged_func = relay.Function([x], out)
merged_func = merged_func.set_attribute('Primitive',
expr.IntImm('int32', 1))
merged_func = merged_func.set_attribute('Composite',
expr.StringImm('C'))
ret = relay.Call(merged_func, [add])
return relay.Function([input_1, input_2], ret)
def after_A():
inputs = [
relay.var('input_' + str(i), shape=(10, 10)) for i in range(4)
]
x = relay.var('x')
y = relay.var('y')
add_relu_1 = relay.add(x, y)
add_relu_1 = relay.nn.relu(add_relu_1)
add_relu_1 = relay.Function([x, y], add_relu_1)
add_relu_1 = add_relu_1.set_attribute('Primitive',
expr.IntImm('int32', 1))
add_relu_1 = add_relu_1.set_attribute('Composite',
expr.StringImm('add_relu'))
add_relu_call_1 = relay.Call(add_relu_1, [inputs[0], inputs[1]])
x1 = relay.var('x1')
y1 = relay.var('y1')
add_relu_2 = relay.add(x1, y1)
add_relu_2 = relay.nn.relu(add_relu_2)
add_relu_2 = relay.Function([x1, y1], add_relu_2)
add_relu_2 = add_relu_2.set_attribute('Primitive',
expr.IntImm('int32', 1))
add_relu_2 = add_relu_2.set_attribute('Composite',
expr.StringImm('add_relu'))
add_relu_call_2 = relay.Call(add_relu_2, [inputs[2], inputs[3]])
x2 = relay.var('x2')
y2 = relay.var('y2')
add = relay.add(x2, y2)
sub = relay.subtract(x2, y2)
add_sub_mul = relay.multiply(add, sub)
add_sub_mul = relay.Function([x2, y2], add_sub_mul)
add_sub_mul = add_sub_mul.set_attribute('Primitive',
expr.IntImm('int32', 1))
add_sub_mul = add_sub_mul.set_attribute('Composite',
expr.StringImm('add_sub_mul'))
add_sub_mul_call = relay.Call(add_sub_mul,
[add_relu_call_1, add_relu_call_2])
return relay.Function(inputs, add_sub_mul_call)
def after_B_priority():
input_1 = relay.var('input_1', shape=(10, 10))
input_2 = relay.var('input_2', shape=(10, 10))
x = relay.var('x')
y = relay.var('y')
out = relay.add(x, y)
out = relay.abs(out)
merged_func = relay.Function([x, y], out)
merged_func = merged_func.set_attribute('Primitive',
expr.IntImm('int32', 1))
merged_func = merged_func.set_attribute('Composite',
expr.StringImm('B'))
merged_call = relay.Call(merged_func, [input_1, input_2])
ret = relay.nn.relu(merged_call)
return relay.Function([input_1, input_2], ret)
def after_B():
inputs = [
relay.var('input_' + str(i), shape=(10, 10)) for i in range(8)
]
add_relu_calls = []
for i in range(4):
x = relay.var('x' + str(i))
y = relay.var('x' + str(i))
add_relu = relay.add(x, y)
add_relu = relay.nn.relu(add_relu)
add_relu = relay.Function([x, y], add_relu)
add_relu = add_relu.set_attribute('Primitive',
expr.IntImm('int32', 1))
add_relu = add_relu.set_attribute('Composite',
expr.StringImm('add_relu'))
add_relu_call = relay.Call(add_relu,
[inputs[i * 2], inputs[i * 2 + 1]])
add_relu_calls.append(add_relu_call)
add = relay.add(add_relu_calls[0], add_relu_calls[1])
sub = relay.subtract(add_relu_calls[2], add_relu_calls[3])
out = relay.multiply(add, sub)
return relay.Function(inputs, out)
