def test_extract_constants_multi():
def _get_func():
var_input1 = relay.var("data1", shape=(10, 10), dtype="uint8")
var_input2 = relay.var("data2", shape=(10, 10), dtype="uint8")
const_data_1 = np.random.uniform(0, 255, (10, 10)).astype("uint8")
const_data_2 = np.random.uniform(0, 255, (10, 10)).astype("uint8")
const_data_3 = np.random.uniform(0, 255, (10, 10)).astype("uint8")
const_data_4 = np.random.uniform(0, 255, (10, 10)).astype("uint8")
const_input_1 = relay.const(const_data_1, dtype="uint8")
const_input_2 = relay.const(const_data_2, dtype="uint8")
const_input_3 = relay.const(const_data_3, dtype="uint8")
const_input_4 = relay.const(const_data_4, dtype="uint8")
out = relay.add(var_input1, var_input2)
out = relay.add(out, const_input_1)
out = relay.add(out, const_input_2)
out = relay.add(out, const_input_3)
out = relay.add(out, const_input_4)
func = relay.Function(relay.analysis.free_vars(out), out)
func = run_opt_pass(func, relay.transform.InferType())
return func, [
const_input_1, const_input_2, const_input_3, const_input_4
]
def _expected():
var_input1 = relay.var("data1", shape=(10, 10), dtype="uint8")
var_input2 = relay.var("data2", shape=(10, 10), dtype="uint8")
var_input3 = relay.var("p1", shape=(10, 10), dtype="uint8")
var_input4 = relay.var("p2", shape=(10, 10), dtype="uint8")
var_input5 = relay.var("p3", shape=(10, 10), dtype="uint8")
var_input6 = relay.var("p4", shape=(10, 10), dtype="uint8")
out = relay.add(var_input1, var_input2)
out = relay.add(out, var_input3)
out = relay.add(out, var_input4)
out = relay.add(out, var_input5)
out = relay.add(out, var_input6)
func = relay.Function(relay.analysis.free_vars(out), out)
func = run_opt_pass(func, relay.transform.InferType())
return func
func, consts = _get_func()
new_func, const_dict = extract_constants(func)
assert tvm.ir.structural_equal(new_func, _expected())
for i, const in enumerate(consts):
assert i + 2 in const_dict
assert (const_dict[i + 2] == consts[i].data.asnumpy()).all()
def test_copy_constants():
ifm_a = relay.var("IFM_A", shape=(1, 26, 26, 32), dtype="int8")
conv_a = make_ethosu_conv2d(ifm_a, 32, 8, (3, 3), (0, 0), (1, 1), (1, 1))
conv_b = make_ethosu_conv2d(conv_a, 8, 4, (1, 1), (0, 0), (1, 1), (1, 1))
func = relay.Function(relay.analysis.free_vars(conv_b), conv_b)
func = run_opt_pass(func, relay.transform.InferType())
func, const_dict = extract_constants(func)
cached_func = lower_to_te(func)
sch = te.create_schedule([cached_func.outputs[0].op])
planner = copy_constants()
planner(cached_func, const_dict, sch)
assert len(sch.stages) == 23
assert ".global" in sch.stages[6].op.name
assert ".global" in sch.stages[8].op.name
assert ".global" in sch.stages[17].op.name
assert ".global" in sch.stages[19].op.name
def test_copy_luts():
ifm_shape = (1, 33, 33, 11)
ifm = relay.var("IFM", shape=ifm_shape, dtype="int8")
lut = relay.const([i for i in range(256)], dtype="int8")
conv = make_ethosu_conv2d(
ifm, ifm_shape[3], 8, (3, 3), (0, 0), (1, 1), (1, 1), lut=lut, activation="TANH"
)
identity = make_ethosu_identity(conv, lut=lut, activation="TANH")
func = relay.Function(relay.analysis.free_vars(identity), identity)
func = run_opt_pass(func, relay.transform.InferType())
func, const_dict = extract_constants(func)
te_graph = lower_to_te(func)
sch = te.create_schedule([te_graph.outputs[0].op])
copy_constants()(te_graph, const_dict, sch)
copy_luts()(te_graph, const_dict, sch)
assert len(sch.stages) == 17
assert ".global" in sch.stages[5].op.name
assert ".global" in sch.stages[7].op.name
assert ".local" in sch.stages[9].op.name
assert ".local" in sch.stages[10].op.name
def test_extract_constants_single():
def _get_func():
var_input = relay.var("data", shape=(10, 10), dtype="uint8")
const_data = np.random.uniform(0, 255, (10, 10)).astype("uint8")
const_input = relay.const(const_data, dtype="uint8")
out = relay.add(var_input, const_input)
func = relay.Function(relay.analysis.free_vars(out), out)
func = run_opt_pass(func, relay.transform.InferType())
return func, const_input
def _expected():
var_input1 = relay.var("data", shape=(10, 10), dtype="uint8")
var_input2 = relay.var("p1", shape=(10, 10), dtype="uint8")
out = relay.add(var_input1, var_input2)
func = relay.Function(relay.analysis.free_vars(out), out)
func = run_opt_pass(func, relay.transform.InferType())
return func
func, const = _get_func()
new_func, const_dict = extract_constants(func)
assert tvm.ir.structural_equal(new_func, _expected())
assert 1 in const_dict
assert (const_dict[1] == const.data.asnumpy()).all()
def create_te_graph(func):
func, consts = extract_constants(func)
mod = tvm.IRModule.from_expr(func)
func = relay.transform.InferType()(mod)["main"]
te_graph = lower_to_te(func)
return te_graph, consts