def test_ethosu_conv2d():
ifm = relay.var("ifm", shape=(1, 10, 20, 30), dtype="uint8")
weight = relay.var("weight", shape=(40, 3, 3, 30), dtype="uint8")
scale_bias = relay.var("scale_bias", shape=(40, 10), dtype="uint8")
lut = relay.var("lut", shape=(), dtype="uint8")
conv = ethosu_ops.ethosu_conv2d(
ifm,
weight,
scale_bias,
lut,
ifm_scale=0.5,
ifm_zero_point=10,
weight_zero_point=12,
ofm_scale=0.25,
ofm_zero_point=14,
ofm_channels=40,
padding=(1, 1, 1, 1),
kernel_shape=(3, 3),
strides=(1, 1),
dilation=(1, 1),
)
expr = relay.Function(relay.analysis.free_vars(conv), conv)
mod = tvm.IRModule.from_expr(expr)
mod = relay.transform.InferType()(mod)
lowered = lower_to_te(mod["main"])
assert len(lowered.outputs) == 1
assert len(lowered.inputs) == 4
conv2d_compute = Convolution2DCompute.from_output(lowered.outputs[0])
assert conv2d_compute.conv2d.name == "ethosu_conv2d"
input_shapes = set()
for inp in lowered.inputs:
input_shapes.add(tuple([x.value for x in inp.shape]))
assert input_shapes == {(40, 10), (1, 10, 20, 30), (40, 3, 3, 30), ()}
def test_inline_no_ops():
input = relay.var("input", shape=(12, 12), dtype="uint8")
slice = relay.strided_slice(input, [0, 0], [6, 6])
relu1 = relay.nn.relu(slice)
reshape = relay.reshape(relu1, (36, ))
relu2 = relay.nn.relu(reshape)
func = relay.Function(relay.analysis.free_vars(relu2), relu2)
func = run_opt_pass(func, relay.transform.InferType())
cached_func = lower_to_te(func)
sch = te.create_schedule([cached_func.outputs[0].op])
inline_no_ops(cached_func, sch)
reshape_tensor = cached_func.outputs[0].op.input_tensors[0]
slice_tensor = reshape_tensor.op.input_tensors[0].op.input_tensors[0]
assert sch[reshape_tensor].attach_type == AttachType.kInline
assert sch[slice_tensor].attach_type == AttachType.kInline
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 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
