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 _cascader(cached_func, const_dict, sch):
weight = cached_func.inputs[1]
scale_bias = cached_func.inputs[2]
out = cached_func.outputs[0]
conv_compute = Convolution2DCompute.from_output(out)
co = conv_compute.split(sch, 3, 10)
cache_weight = sch.cache_read(weight, "global", [conv_compute.conv2d])
cache_scale_bias = sch.cache_read(scale_bias, "global", [conv_compute.conv2d])
sch[cache_weight].compute_at(sch[out], co)
sch[cache_scale_bias].compute_at(sch[out], co)
def _planner(te_graph, const_dict, sch):
weights = te_graph.inputs[1]
bias = te_graph.inputs[2]
out = te_graph.outputs[0]
conv_compute = Convolution2DCompute.from_output(out)
co = conv_compute.split(sch, 3, 2)
cache_weights = sch.cache_read(weights, "global", [conv_compute.conv2d])
cache_bias = sch.cache_read(bias, "global", [conv_compute.conv2d])
sch[cache_weights].compute_at(sch[out], co)
sch[cache_bias].compute_at(sch[out], co)