def test_tuple_fst(target, dev):
ttype = relay.TupleType([relay.TensorType((1,)), relay.TensorType((10,))])
tup = relay.var("tup", type_annotation=ttype)
f = relay.Function([tup], relay.TupleGetItem(tup, 0))
i_data = np.random.rand(41).astype("float32")
j_data = np.random.rand(10).astype("float32")
mod = tvm.IRModule()
mod["main"] = f
check_result(target, dev, [(i_data, j_data)], i_data, mod=mod)
def test_tuple_getitem(use_calculated_workspaces, target_options):
func = relay.Function([],
relay.TupleGetItem(
relay.Tuple([relay.const(1),
relay.const(2)]), 0))
output_list = generate_ref_data(func, {})
input_list = []
compile_and_run(func, input_list, output_list, target_options,
use_calculated_workspaces)
def test_tuple_output(interface_api, use_unpacked_api, test_runner):
x = relay.var("x", shape=(6, 9))
y = relay.split(x, 3).astuple()
a = relay.TupleGetItem(y, 0)
b = relay.TupleGetItem(y, 1)
out = relay.Tuple([a, b])
func = relay.Function([x], out)
x_data = np.random.rand(6, 9).astype("float32")
inputs = {"x": x_data}
output_list = generate_ref_data(func, inputs)
compile_and_run(
AOTTestModel(module=IRModule.from_expr(func),
inputs=inputs,
outputs=output_list),
test_runner,
interface_api,
use_unpacked_api,
)
def test_square_second_order():
shape = (10, 10)
dtype = 'float32'
t = relay.TensorType(shape, dtype)
x = relay.var("x", t)
func = relay.Function([x], x * x)
func = run_infer_type(func)
back_func = run_infer_type(gradient(func))
y = relay.var("y", t)
back_func_adjusted = relay.Function([y], relay.TupleGetItem(relay.TupleGetItem(back_func(y), 1), 0))
back_func_adjusted = run_infer_type(back_func_adjusted)
back_back_func = run_infer_type(gradient(back_func_adjusted))
assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])]))
x_nd = rand(dtype, *shape)
ex = create_executor()
forward, (grad_x,) = ex.evaluate(back_back_func)(x_nd)
tvm.testing.assert_allclose(forward.asnumpy(), 2 * x_nd.asnumpy())
tvm.testing.assert_allclose(grad_x.asnumpy(), 2 * np.ones_like(grad_x.asnumpy()))
def test_checkpoint():
inputs = [relay.var("x{}".format(i), shape=(1,)) for i in range(4)]
output = relay.multiply(relay.add(inputs[0], inputs[1]), relay.add(inputs[2], inputs[3]))
check_grad(relay.Function(inputs, relay.annotation.checkpoint(output)))
scope = relay.ScopeBuilder()
out_tuple = scope.let(
"out_tuple",
relay.Tuple([relay.add(inputs[0], inputs[1]), relay.multiply(inputs[2], inputs[3])]),
)
scope.ret(
relay.subtract(
relay.annotation.checkpoint(relay.TupleGetItem(out_tuple, 0)),
relay.TupleGetItem(out_tuple, 1),
)
)
out_single = scope.get()
check_grad(relay.Function(inputs, out_single))
def test_split_no_fuse():
x = relay.var('x', shape=(12, ))
y = relay.split(x, 3, axis=0).astuple()
z = relay.concatenate([relay.TupleGetItem(y, 0)], axis=0)
z = relay.annotation.stop_fusion(z)
f = relay.Function([x], z)
x_data = np.random.rand(12, ).astype('float32')
res = veval(f, x_data)
tvm.testing.assert_allclose(res.asnumpy(), np.split(x_data, 3, axis=0)[0])
def test_tuple_second():
ttype = relay.TupleType([relay.TensorType((1,)), relay.TensorType((10,))])
tup = relay.var('tup', type_annotation=ttype)
f = relay.Function([tup], relay.TupleGetItem(tup, 1))
i_data = np.random.rand(41).astype('float32')
j_data = np.random.rand(10).astype('float32')
mod = tvm.IRModule()
mod["main"] = f
check_result([(i_data, j_data)], j_data, mod=mod)
def get_net(iterations, num_hidden, batch_size=1, dtype="float32"):
'''Constructs an unrolled RNN with LSTM cells'''
input_type = relay.TensorType((batch_size, num_hidden), dtype)
weight_type = relay.TensorType((4 * num_hidden, num_hidden), dtype)
bias_type = relay.TensorType((4 * num_hidden, ), dtype)
state_type = relay.TupleType([input_type, input_type])
cell_type = relay.TupleType([input_type, state_type])
builder = relay.ScopeBuilder()
zeros = builder.let(("zeros", input_type),
relay.zeros((batch_size, num_hidden), dtype))
init_states = builder.let(("init_states", state_type),
relay.Tuple([zeros, zeros]))
states = init_states
out = None
for i in range(iterations):
inputs = relay.Var("data", input_type)
i2h_weight = relay.Var("i2h_%s_weight" % i, weight_type)
i2h_bias = relay.Var("i2h_%i_bias" % i, bias_type)
h2h_weight = relay.Var("h2h_%s_weight" % i, weight_type)
h2h_bias = relay.Var("h2h_%s_bias" % i, bias_type)
cell_fn = lstm_cell(num_hidden, batch_size, dtype, "lstm_%s" % i)
call = builder.let(
("call_%s" % i, cell_type),
relay.Call(
cell_fn,
[inputs, states, i2h_weight, i2h_bias, h2h_weight, h2h_bias]))
new_out = builder.let(("out_%s" % i, input_type),
relay.TupleGetItem(call, 0))
new_states = builder.let(("states_%s" % i, state_type),
relay.TupleGetItem(call, 1))
states = new_states
out = new_out
builder.ret(out)
body = builder.get()
args = relay.analysis.free_vars(body)
return relay.Function(args, body, input_type)
def expected():
mod = tvm.IRModule()
# function 1
f1_cb1 = relay.var('test_target_1_i0', shape=(10, 10))
f1_O_1 = relay.abs(f1_cb1)
f1_O_2 = relay.nn.relu(f1_O_1)
f1_out = relay.Tuple((f1_O_2, f1_O_1))
func1 = relay.Function([f1_cb1], f1_out)
func1 = func1.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func1 = func1.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func1 = func1.with_attr("Compiler", tvm.tir.StringImm("test_target"))
func1 = func1.with_attr("global_symbol",
container.String("test_target_1"))
gv1 = relay.GlobalVar("test_target_1")
mod[gv1] = func1
# function 0
f2_cb3 = relay.var('test_target_0_i0', shape=(10, 10))
f2_cb4 = relay.var('test_target_0_i1', shape=(10, 10))
f2_O_3 = relay.add(f2_cb3, f2_cb4)
func0 = relay.Function([f2_cb3, f2_cb4], f2_O_3)
func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Compiler", tvm.tir.StringImm("test_target"))
func0 = func0.with_attr("global_symbol",
container.String("test_target_0"))
gv0 = relay.GlobalVar("test_target_0")
mod[gv0] = func0
# body
data = relay.var('data', shape=(10, 10))
tuple_out = gv1(data)
ce_2 = relay.TupleGetItem(tuple_out, 1)
ce_3 = relay.TupleGetItem(tuple_out, 0)
X = relay.tanh(ce_2)
ce_4 = gv0(ce_3, X)
func = relay.Function([data], ce_4)
mod["main"] = func
return mod
def get_model(get_item):
"""Return a model"""
a = relay.var("a", shape=(1, 16, 16, 4), dtype="uint8")
z = relay.op.clip(a, 0, 255)
b = relay.op.clip(z, 0, 15)
c = relay.op.clip(z, 16, 31)
t = relay.Tuple((c, b))
tgi = relay.TupleGetItem(t, 1) if get_item else t
foo = relay.Function([a], tgi)
return tvm.IRModule.from_expr(tgi)
def test_tuple_second():
ttype = relay.TupleType(
[relay.TensorType((1, )),
relay.TensorType((10, ))])
tup = relay.var('tup', type_annotation=ttype)
f = relay.Function([tup], relay.TupleGetItem(tup, 1))
i_data = np.random.rand(41).astype('float32')
j_data = np.random.rand(10).astype('float32')
result = veval(f, (i_data, j_data))
tvm.testing.assert_allclose(result.asnumpy(), j_data)
def relay_lstm_cell(batch_size, input_size, hidden_size):
# based on https://pytorch.org/docs/stable/generated/torch.nn.GRU.html#torch.nn.GRU
state_tensor_type = relay.TensorType((batch_size, hidden_size))
state_tuple_type = relay.TupleType([state_tensor_type, state_tensor_type])
inp = relay.var("input", shape=(batch_size, input_size))
state = relay.Var("state", type_annotation=state_tuple_type)
w_ih = relay.var("w_ih", shape=(4*hidden_size, input_size))
w_hh = relay.var("w_hh", shape=(4*hidden_size, hidden_size))
b_ih = relay.var("b_ih", shape=(4*hidden_size,))
b_hh = relay.var("b_hh", shape=(4*hidden_size,))
hidden = relay.TupleGetItem(state, 0)
cell_state = relay.TupleGetItem(state, 1)
# PyTorch packs the i2h and h2h weights and biases together so we will match that here
w_i_splits = relay.split(w_ih, 4, 0)
w_h_splits = relay.split(w_hh, 4, 0)
b_i_splits = relay.split(b_ih, 4, 0)
b_h_splits = relay.split(b_hh, 4, 0)
w_ii, w_if, w_ig, w_io = w_i_splits[0], w_i_splits[1], w_i_splits[2], w_i_splits[3]
w_hi, w_hf, w_hg, w_ho = w_h_splits[0], w_h_splits[1], w_h_splits[2], w_h_splits[3]
b_ii, b_if, b_ig, b_io = b_i_splits[0], b_i_splits[1], b_i_splits[2], b_i_splits[3]
b_hi, b_hf, b_hg, b_ho = b_h_splits[0], b_h_splits[1], b_h_splits[2], b_h_splits[3]
def weighted_value(weight, value, bias):
return relay.transpose(relay.nn.dense(weight, value) + relay.reshape(bias, (hidden_size, 1)))
i_t = relay.sigmoid(weighted_value(w_ii, inp, b_ii) + weighted_value(w_hi, hidden, b_hi))
f_t = relay.sigmoid(weighted_value(w_if, inp, b_if) + weighted_value(w_hf, hidden, b_hf))
g_t = relay.tanh(weighted_value(w_ig, inp, b_ig) + weighted_value(w_hg, hidden, b_hg))
o_t = relay.sigmoid(weighted_value(w_io, inp, b_io) + weighted_value(w_ho, hidden, b_ho))
c_t = f_t*cell_state + i_t*g_t
h_t = o_t*relay.tanh(c_t)
h_var = relay.Var("h")
c_var = relay.Var("c")
return relay.Function([inp, state, w_ih, w_hh, b_ih, b_hh],
relay.Let(h_var, h_t,
relay.Let(c_var, c_t,
relay.Tuple([h_var, relay.Tuple([h_var, c_var])]))),
ret_type=relay.TupleType([state_tensor_type, state_tuple_type]))
def test_tuple_output(interface_api, use_unpacked_api,
use_calculated_workspaces):
x = relay.var("x", shape=(6, 9))
y = relay.split(x, 3).astuple()
a = relay.TupleGetItem(y, 0)
b = relay.TupleGetItem(y, 1)
out = relay.Tuple([a, b])
func = relay.Function([x], out)
x_data = np.random.rand(6, 9).astype("float32")
inputs = {"x": x_data}
output_list = generate_ref_data(func, inputs)
compile_and_run(
func,
inputs,
output_list,
interface_api,
use_unpacked_api,
use_calculated_workspaces,
)
def test_tuple_getitem(interface_api, use_unpacked_api, test_runner):
func = relay.Function([], relay.TupleGetItem(relay.Tuple([relay.const(1), relay.const(2)]), 0))
output_list = generate_ref_data(func, {})
compile_and_run(
AOTTestModel(module=IRModule.from_expr(func), inputs={}, outputs=output_list),
test_runner,
interface_api,
use_unpacked_api,
)
def test_split_no_fuse():
x = relay.var("x", shape=(12, ))
y = relay.split(x, 3, axis=0).astuple()
z = relay.concatenate([relay.TupleGetItem(y, 0)], axis=0)
z = relay.annotation.stop_fusion(z)
f = relay.Function([x], z)
x_data = np.random.rand(12, ).astype("float32")
for tgt, dev in tvm.testing.enabled_targets():
res = veval(f, x_data, device=dev, target=tgt)
tvm.testing.assert_allclose(res.numpy(),
np.split(x_data, 3, axis=0)[0])
def test_tuple():
ttype = relay.TupleType(
[relay.TensorType((1, )),
relay.TensorType((10, ))])
tup = relay.var("tup", type_annotation=ttype)
f = relay.Function([tup], relay.TupleGetItem(tup, 1))
i_data = np.random.rand(41).astype("float32")
j_data = np.random.rand(10).astype("float32")
result = get_serialized_output(f, (i_data, j_data))
tvm.testing.assert_allclose(result.asnumpy(), j_data)
def create_graph(axis, sections):
x = relay.var("x", shape=(1, 50, 50, 3))
x_abs = relay.abs(x)
split_output = relay.split(x_abs, sections, axis).tuple_value
outputs = list()
for section_idx in range(sections):
split_single_out = relay.TupleGetItem(split_output, section_idx)
tanh = relay.tanh(split_single_out)
outputs.append(tanh)
tuple_out = relay.Tuple(outputs)
return relay.Function([x], tuple_out)
def test_tuple():
mod = Module()
cfunc = compile(
Function([],
relay.TupleGetItem(
relay.Tuple([
relay.const(3, dtype='int32'),
relay.const(4.0, dtype='float32')
]), 1)), mod)
np.testing.assert_allclose(cfunc().asnumpy(), np.array(4.0,
dtype='float32'))
def _test_tuple_argument(mode):
shape = (2, 3)
dtype = "float32"
tensor_type = relay.TensorType(shape, dtype)
fields = 3
tuple_type = relay.TupleType([tensor_type] * fields)
tup = relay.var("tup", type_annotation=tuple_type)
body = relay.TupleGetItem(tup, 0)
for i in range(1, fields):
body = relay.add(body, relay.TupleGetItem(tup, i))
func = relay.Function([tup], body)
func = run_infer_type(func)
back_func = run_infer_type(gradient(func, mode=mode))
xs = [rand(dtype, *shape) for _ in range(fields)]
xs_np = np.array([x.numpy() for x in xs])
expected_forward = np.sum(xs_np, axis=0)
forward, grad = create_executor().evaluate(back_func)(tuple(xs))
tvm.testing.assert_allclose(forward.numpy(), expected_forward)
for field in grad[0]:
tvm.testing.assert_allclose(field.numpy(), np.ones_like(field.numpy()))
def build_relay_module(batch_size, input_size, hidden_size, time_steps,
dense_dim):
mod = tvm.IRModule()
mod["lstm_layer"] = lstm_definition(batch_size, input_size, hidden_size,
time_steps)
mod["linear_layer"] = linear_layer_definition(batch_size, hidden_size,
dense_dim)
lstm_var = mod.get_global_var("lstm_layer")
linear_var = mod.get_global_var("linear_layer")
# now we build up our main function
input_var = relay.var("input", shape=(batch_size, time_steps, input_size))
init_hidden_var = relay.var("init_hidden", shape=(batch_size, hidden_size))
init_cell_var = relay.var("init_cell", shape=(batch_size, hidden_size))
i2h_weight_var = relay.var("i2h_weight",
shape=(4 * hidden_size, input_size))
h2h_weight_var = relay.var("h2h_weight",
shape=(4 * hidden_size, hidden_size))
lstm_bias_var = relay.var("lstm_bias", shape=(4 * hidden_size, ))
linear_weight_var = relay.var("linear_weight",
shape=(dense_dim, hidden_size))
linear_bias_var = relay.var("linear_bias", shape=(dense_dim, ))
builder = relay.ScopeBuilder()
state_var = builder.let("state",
relay.Tuple([init_hidden_var, init_cell_var]))
lstm_res = builder.let(
"lstm_res",
lstm_var(
input_var,
state_var,
i2h_weight_var,
h2h_weight_var,
lstm_bias_var,
# the keras model only gave one bias,
# so set the other to zero
# (hopefully this is correct)
relay.zeros_like(lstm_bias_var)))
final_hidden = builder.let("final_hidden", relay.TupleGetItem(lstm_res, 1))
# to match PT's semantics, we're undoing the reshape in LSTM :)
reshape_hidden = builder.let("reshape_hidden",
relay.squeeze(final_hidden, axis=[0]))
linear_result = builder.let(
"linear_result",
linear_var(reshape_hidden, linear_weight_var, linear_bias_var))
# finally do a softmax
builder.ret(relay.nn.softmax(linear_result))
main_func = relay.Function([
input_var, init_hidden_var, init_cell_var, i2h_weight_var,
h2h_weight_var, lstm_bias_var, linear_weight_var, linear_bias_var
], builder.get())
mod["main"] = main_func
return mod
def test_compile_fused_identity_cast():
# a fused function that would optimized to identity
x = relay.var("x", shape=[16], dtype="float32")
y = relay.cast(x, "float32")
func1 = relay.Function([x], y).with_attr("Primitive", 1)
# a fused function with param pass-through
x = relay.var("x", shape=[16], dtype="float32")
y = relay.add(x, relay.const(3.14, "float32"))
func2 = relay.Function([x], relay.Tuple([x, y])).with_attr("Primitive", 1)
x_global = relay.var("xx", shape=[16], dtype="float32")
tup = func2(x_global)
y_global = func1(relay.TupleGetItem(tup, 0) + relay.TupleGetItem(tup, 1))
mod = tvm.IRModule.from_expr(relay.Function([x_global], y_global))
for target, device in tvm.testing.enabled_targets():
with tvm.transform.PassContext(opt_level=2):
graph, lib, _ = relay.build(mod, target=target)
executor = graph_executor.create(graph, lib, device=device)
executor.run()
def expected():
mod = tvm.IRModule()
# function 0
f0_i0 = relay.var(target + "_0_i0", shape=(10, 10), dtype="uint8")
a_split = relay.split(f0_i0, 2)
a_split_0 = relay.TupleGetItem(a_split.astuple(), 0)
a_split_1 = relay.TupleGetItem(a_split.astuple(), 1)
a_split_abs_in = relay.TupleGetItem(a_split.astuple(), 0)
abs = relay.abs(a_split_abs_in)
tuple_out = relay.Tuple((a_split_0, a_split_1, abs))
func0 = relay.Function([f0_i0], tuple_out)
func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Compiler", target)
func0 = func0.with_attr("global_symbol", target + "_0")
gv0 = relay.GlobalVar(target + "_0")
mod[gv0] = func0
#body
data = relay.var('a', shape=(10, 10), dtype="uint8")
f_out = gv0(data)
f_out_0 = relay.TupleGetItem(f_out, 0)
f_out_1 = relay.TupleGetItem(f_out, 1)
tuple = relay.Tuple((f_out_0, f_out_1))
concat = relay.concatenate(tuple, 0)
f_out_2 = relay.TupleGetItem(f_out, 2)
relu = relay.nn.relu(f_out_2)
ret_tuple = relay.Tuple((concat, relu))
mod["main"] = relay.Function([data], ret_tuple)
return mod
def expected():
def create_external_func1(mod_, compiler_name, symbol_name):
x_int = relay.var("x_int", shape=(10, 10))
p0 = relay.nn.relu(x_int)
q0 = relay.tanh(x_int)
# reshapes
p0_reshaped = relay.reshape(p0, newshape=100)
q0_reshaped = relay.reshape(q0, newshape=100)
ofms = relay.concatenate((p0_reshaped, q0_reshaped), 0)
f1 = relay.Function([x_int], ofms)
f1 = set_func_attr(f1, compiler_name, symbol_name)
glb_f1 = relay.GlobalVar(symbol_name)
mod_[glb_f1] = f1
mod_ = relay.transform.InferType()(mod_)
return glb_f1, mod_
mod = tvm.IRModule()
x = relay.var("x", shape=(10, 10))
glb_symbol_f1, mod = create_external_func1(mod, "ethosu", "ethosu_0")
ofms = relay.Call(glb_symbol_f1, [x])
# splits
(p0_flat, q0_flat) = relay.split(ofms, [100])
# reshapes
p0_flat_reshaped = relay.reshape(p0_flat, newshape=(10, 10))
q0_flat_reshaped = relay.reshape(q0_flat, newshape=(10, 10))
# original output
tuple_out = relay.Tuple([p0_flat_reshaped, q0_flat_reshaped])
p0 = relay.TupleGetItem(tuple_out, 0)
q0 = relay.TupleGetItem(tuple_out, 1)
r = relay.concatenate((p0, q0), axis=0)
main = relay.Function([x], r)
mod["main"] = main
mod = relay.transform.InferType()(mod)
return mod
def test_vm_reshape_tuple(target, dev, x_shape=(1, 4, 2), y_shape=(1, 2, 10)):
tup = relay.var(
"tup",
type_annotation=relay.TupleType([relay.TensorType(x_shape), relay.TensorType(y_shape)]),
)
out = relay.reshape(relay.TupleGetItem(tup, 0), (1, -1))
f = relay.Function([tup], out)
x_data = np.random.uniform(size=x_shape).astype("float32")
y_data = np.random.uniform(size=y_shape).astype("float32")
res = veval(f, (x_data, y_data), device=dev, target=target)
tvm.testing.assert_allclose(res.numpy(), np.reshape(x_data, (1, -1)))
def expected():
mod = tvm.IRModule()
# function 0
f0_i0 = relay.var(target + "_0_i0", shape=(10, 10))
f0_i1 = relay.var(target + "_0_i1")
f0_i2 = relay.var(target + "_0_i2")
f0_i3 = relay.var(target + "_0_i3")
f0_i4 = relay.var(target + "_0_i4")
f0_n0 = relay.nn.batch_norm(f0_i0, f0_i1, f0_i2, f0_i3, f0_i4)
f0_n1 = f0_n0[1]
f0_n2 = relay.nn.relu(f0_n0[0])
f0_o0 = relay.Tuple([f0_n2, f0_n1])
func0 = relay.Function([f0_i0, f0_i1, f0_i2, f0_i3, f0_i4], f0_o0)
func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
func0 = func0.with_attr("Compiler", target)
func0 = func0.with_attr("global_symbol", target + "_0")
gv0 = relay.GlobalVar(target + "_0")
mod[gv0] = func0
mod = transform.InferType()(mod)
# body
data = relay.var("data", shape=(10, 10))
bn_gamma = relay.var("bn_gamma")
bn_beta = relay.var("bn_beta")
bn_mmean = relay.var("bn_mean")
bn_mvar = relay.var("bn_var")
function_out = gv0(data, bn_gamma, bn_beta, bn_mmean, bn_mvar)
get_out0 = relay.TupleGetItem(function_out, 0)
get_out1 = relay.TupleGetItem(function_out, 1)
out_2 = relay.tanh(get_out1)
out_3 = relay.log(get_out1)
out = relay.Tuple([get_out0, out_2, out_3])
func = relay.Function([data, bn_gamma, bn_beta, bn_mmean, bn_mvar], out)
mod["main"] = func
mod = transform.InferType()(mod)
return mod
def do_env_find(self, env, key, dft):
"""Build the code to find a value in env."""
v = relay.var("v")
cl = adt.Clause(
adt.PatternConstructor(self.env_ctr, [adt.PatternVar(v)]), v)
env_v = adt.Match(env, [cl], complete=False)
val = relay.TupleGetItem(env_v, self.env_val_map[key][0])
x = relay.var("x")
nil_c = adt.Clause(adt.PatternConstructor(nil, []), dft)
some_c = adt.Clause(adt.PatternConstructor(some, [adt.PatternVar(x)]),
x)
return adt.Match(val, [some_c, nil_c])
def do_env_update(self, env_, key, val):
"""Build the code to update the env."""
v = relay.var("v")
cl = adt.Clause(
adt.PatternConstructor(self.env_ctr, [adt.PatternVar(v)]), v)
env = adt.Match(env_, [cl], complete=False)
map = dict((i, k) for k, (i, _) in self.env_val_map.items())
new_env = relay.Tuple([
some(val) if map[i] == key else relay.TupleGetItem(env, i)
for i in range(len(map))
])
return self.env_ctr(new_env)
def fuse_partitions(pre_mod, mid_mod, post_mod):
"""Combine prefix, middle, and suffix modules into a single module.
The combined module includes an additional `main` that fuses all three
partitions together.
Parameters
----------
pre_mod : tvm.IRModule
Module containing an input quantization function
mid_mod : tvm.IRModule
Module containing core of a quantized inference function
post_mod : tvm.IRModule
Module containing an output dequantization function
Returns
-------
fused_mod : tvm.IRModule
Module containing the input quantization, core quantized inference,
output dequantization, and full quantized inference functions
"""
pre_func = pre_mod['main']
mid_func = mid_mod['main']
post_func = post_mod['main']
# create a module containing the prefix, middle, and suffix partitions
fused_mod = tvm.IRModule(functions={
relay.GlobalVar('quantize_inputs'): pre_func,
relay.GlobalVar('quantized_main'): mid_func,
relay.GlobalVar('dequantize_outputs'): post_func,
})
# construct a `main` that strings together the partitions, such that its
# behaviour is equivalent to `main` in an *unpartitioned* module
scope_builder = relay.ScopeBuilder()
fused_mod_main_params = [relay.Var(param.name_hint) for param in pre_func.params]
quantized_inputs = scope_builder.let('quantized_inputs', relay.Call(
fused_mod.get_global_var('quantize_inputs'),
fused_mod_main_params
))
quantized_outputs = scope_builder.let('quantized_outputs', relay.Call(
fused_mod.get_global_var('quantized_main'),
[relay.TupleGetItem(quantized_inputs, i) for i in range(len(pre_func.ret_type.fields))]
))
dequantized_outputs = scope_builder.let('dequantized_outputs', relay.Call(
fused_mod.get_global_var('dequantize_outputs'),
[quantized_outputs]
))
scope_builder.ret(dequantized_outputs)
fused_mod['main'] = relay.Function(fused_mod_main_params, scope_builder.get())
return fused_mod
def test_simple_graph():
# A module with two subgraphs
mod = tvm.IRModule()
x0 = relay.var("x0", shape=(8, 8))
y0 = relay.var("y0", shape=(8, 8))
z0 = x0 + y0
z1 = x0 - y0
z2 = relay.Tuple((z0, z1))
f0 = relay.Function([x0, y0], z2)
f0 = f0.with_attr("Compiler", "test_graph")
g0 = relay.GlobalVar("g0")
mod[g0] = f0
mod = relay.transform.InferType()(mod)
x1 = relay.var("x1", shape=(8, 8))
y1 = relay.var("y1", shape=(8, 8))
z1 = x1 - y1
f1 = relay.Function([x1, y1], z1)
f1 = f1.with_attr("Compiler", "test_graph")
g1 = relay.GlobalVar("g1")
mod[g1] = f1
mod = relay.transform.InferType()(mod)
x = relay.var("x", shape=(8, 8))
y = relay.var("y", shape=(8, 8))
z = relay.var("z", shape=(8, 8))
c0 = relay.Call(g0, [x, y])
c1 = relay.Call(g1, [relay.TupleGetItem(c0, 0), z])
fm = relay.Function([x, y, z], c1)
mod["main"] = fm
mod = relay.transform.InferType()(mod)
x_data = np.random.rand(8, 8).astype("float32")
y_data = np.random.rand(8, 8).astype("float32")
z_data = np.random.rand(8, 8).astype("float32")
data = get_calibration_data(mod, {"x": x_data, "y": y_data, "z": z_data})
# Check the number and orders
check_data_size(mod, data)
tvm.testing.assert_allclose(data[g0]["inputs"][0].asnumpy(), x_data)
tvm.testing.assert_allclose(data[g0]["inputs"][1].asnumpy(), y_data)
tvm.testing.assert_allclose(data[g0]["outputs"][0].asnumpy(),
x_data + y_data)
tvm.testing.assert_allclose(data[g0]["outputs"][1].asnumpy(),
x_data - y_data)
tvm.testing.assert_allclose(data[g1]["inputs"][0].asnumpy(),
x_data + y_data)
tvm.testing.assert_allclose(data[g1]["inputs"][1].asnumpy(), z_data)
tvm.testing.assert_allclose(data[g1]["outputs"][0].asnumpy(),
x_data + y_data - z_data)
def test_tuple_get_item_sequal():
x = relay.Var('x')
y = relay.Var('y')
assert not consistent_equal(relay.TupleGetItem(x, 1),
relay.TupleGetItem(y, 1))
assert not consistent_equal(relay.TupleGetItem(x, 1),
relay.TupleGetItem(x, 2))
assert consistent_equal(relay.TupleGetItem(x, 1), relay.TupleGetItem(x, 1))
