def test_conv2d_inline_reshape(trial):
def _get_func(ifm_shape, reshaped, ifm_layout):
ifm = relay.var("ifm", shape=ifm_shape, dtype="int8")
ifm_reshaped = relay.reshape(ifm, reshaped)
conv = make_ethosu_conv2d(
ifm_reshaped,
reshaped[3],
16,
(3, 3),
(1, 1),
(1, 1),
(1, 1),
activation="NONE",
ifm_layout=ifm_layout,
)
func = relay.Function(relay.analysis.free_vars(conv), conv)
func = run_opt_pass(func, relay.transform.InferType())
return func
reference_mod = trial[0]
params = trial[1:]
func = _get_func(*params)
mod, _ = _lower_to_tir(func, cascader=total_cascader((1, 4, 6, 16)))
script = mod.script(show_meta=True)
mod = tvm.script.from_source(script)
tvm.ir.assert_structural_equal(mod["main"], reference_mod["main"], True)
def test_conv2d_big_pad():
def _get_func():
ifm_shape = (1, 2, 2, 8)
ifm = relay.var("ifm", shape=ifm_shape, dtype="int8")
conv = make_ethosu_conv2d(ifm, ifm_shape[3], 16, (1, 1), (7, 7), (1, 1), (1, 1), "NHWC")
func = relay.Function(relay.analysis.free_vars(conv), conv)
func = run_opt_pass(func, relay.transform.InferType())
return func
func = _get_func()
mod, _ = lower_to_tir(func, cascader=total_cascader((1, 4, 4, 16)))
def test_conv2d_double_cascade(trial):
def _get_func(
ifm_shape,
ifm_channels,
mid_channels,
ofm_channels,
kernel_shape,
padding,
strides,
dilation,
layout,
upscale,
):
ifm = relay.var("ifm", shape=ifm_shape, dtype="int8")
conv1 = make_ethosu_conv2d(
ifm,
ifm_channels,
mid_channels,
kernel_shape,
padding,
strides,
dilation,
activation="NONE",
ifm_layout=layout,
ofm_layout=layout,
upscale=upscale,
)
conv2 = make_ethosu_conv2d(
conv1,
mid_channels,
ofm_channels,
kernel_shape,
padding,
strides,
dilation,
activation="NONE",
ifm_layout=layout,
ofm_layout=layout,
upscale=upscale,
)
func = relay.Function(relay.analysis.free_vars(conv2), conv2)
func = run_opt_pass(func, relay.transform.InferType())
return func
reference_mod = trial[0]
params = trial[1:]
func = _get_func(*params[:-1])
mod, _ = _lower_to_tir(func, cascader=total_cascader(params[-1]))
script = mod.script(show_meta=True)
mod = tvm.script.from_source(script)
tvm.ir.assert_structural_equal(mod["main"], reference_mod["main"], True)
def test_total_cascader():
input = te.placeholder((12, 12), dtype="uint8", name="input")
relu1 = topi.nn.relu(input)
relu2 = topi.nn.relu(relu1)
relu3 = topi.nn.relu(relu2)
sch = te.create_schedule([relu3.op])
cascader = total_cascader((4, 4))
cascader(TestTEGraph([input], [relu3]), {}, sch)
assert sch[relu1].attach_type == AttachType.kScope
assert sch[relu2].attach_type == AttachType.kScope
assert sch[relu3].attach_type == AttachType.kGroupRoot
# Check that the attaches are at the correct iter var
assert sch[relu1].attach_ivar == sch[relu3].leaf_iter_vars[1]
assert sch[relu2].attach_ivar == sch[relu3].leaf_iter_vars[1]
