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)