| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| """Test the layout optimization pass. This pass is used to |
| convert subgraphs to the preferred layout of NHCWB16. |
| """ |
| |
| import pytest |
| |
| pytest.importorskip("ethosu.vela") |
| |
| import sys |
| |
| import numpy as np |
| import tensorflow as tf |
| import tflite.Model |
| |
| import tvm |
| from tvm import relay |
| from tvm.relay.op.contrib.ethosu import partition_for_ethosu |
| from tvm.relay.backend.contrib.ethosu.codegen import LayoutOptimizer |
| from tvm.relay.backend.contrib.ethosu.codegen import relay_to_tir |
| |
| from . import infra |
| |
| |
| def _optimize(func, optimize=True): |
| """Create IRModule and run layout optimizer pass.""" |
| func = func.with_attr("Compiler", "ethos-u") |
| mod = tvm.IRModule.from_expr(func) |
| mod = relay.transform.InferType()(mod) |
| if optimize: |
| mod = LayoutOptimizer()(mod) |
| entry = mod["main"] |
| return entry if isinstance(func, relay.Function) else entry.body |
| |
| |
| def _assert_structural_equal(a, b): |
| """Check structural equality of two Relay expressions.""" |
| reason = ( |
| "Actual and expected relay functions are not equal. " |
| "LayoutOptimizer is not correctly converting layouts." |
| ) |
| assert tvm.ir.structural_equal(a, b), reason |
| |
| |
| def _compile_and_compare_model(tflite_graph, ifm_shape, dtype): |
| """Compare running result of compilation against TFLite.""" |
| tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0) |
| |
| mod, params = relay.frontend.from_tflite( |
| tflite_model, |
| shape_dict={ |
| "ifm": ifm_shape, |
| }, |
| dtype_dict={ |
| "ifm": dtype, |
| }, |
| ) |
| mod = partition_for_ethosu(mod, params) |
| |
| # Generate reference data |
| input_data, output_data = infra.generate_ref_data_tflite(tflite_graph) |
| |
| test_runner = infra.create_test_runner("ethos-u55-256") |
| compiled_models = infra.build_source( |
| mod, |
| input_data, |
| output_data, |
| test_runner, |
| output_tolerance=0, |
| ) |
| |
| # Assumes only two runtime.Modules are created -- i.e. single offload module |
| ethosu_module = compiled_models[0].executor_factory.lib.imported_modules[0].imported_modules[0] |
| |
| # Verify generated C source |
| get_artifacts = tvm._ffi.get_global_func("runtime.module.ethos-u.get_artifacts") |
| compilation_artifacts = get_artifacts(ethosu_module) |
| cmms = bytes.fromhex(compilation_artifacts[0].command_stream) |
| infra.print_payload(cmms) |
| infra.verify_source(compiled_models, test_runner) |
| |
| |
| def test_single_convolution(): |
| """Test a single convolution to make sure the layouts remain |
| unaltered. |
| """ |
| |
| def get_graph(): |
| x = relay.var("x", shape=(1, 8, 8, 4), dtype="int8") |
| x = infra.make_ethosu_conv2d( |
| ifm=x, |
| ifm_channels=8, |
| ofm_channels=8, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(x), x) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| @pytest.mark.parametrize("dtype", ["int8", "int32"]) |
| def test_add_reduce_sum(dtype): |
| """Test add with reduce sum to make sure the layouts remain |
| unaltered for int32 and altered for other types. |
| """ |
| |
| def get_graph(get_expected=False): |
| in_1 = relay.var("x", shape=(1, 2, 2, 2), dtype=dtype) |
| in_2 = relay.var("y", shape=(1, 2, 2, 2), dtype=dtype) |
| layout = "NHCWB16" if get_expected and dtype != "int32" else "NHWC" |
| add = infra.make_ethosu_binary_elementwise( |
| in_1, |
| in_2, |
| ifm_channels=2, |
| ifm2_channels=2, |
| operator_type="ADD", |
| ofm_dtype=dtype, |
| ifm_layout="NHWC", |
| ifm2_layout="NHWC", |
| ofm_layout=layout, |
| ) |
| x = infra.make_ethosu_pooling( |
| ifm=add, |
| pooling_type="SUM", |
| pool_shape=(1, 1), |
| ofm_channels=1, |
| ofm_dtype="int32", |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout=layout, |
| ofm_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(x), x) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_multiple_convolution(): |
| """Test layout optimization pass on linear chain of convolutions. I.e, |
| |
| conv_1 |
| | |
| conv_2 |
| | |
| conv_3 |
| """ |
| |
| def get_graph(get_expected=False): |
| x = relay.var("x", shape=(1, 8, 8, 4), dtype="int8") |
| for i in range(3): |
| ifm_layout = "NHCWB16" if get_expected and i != 0 else "NHWC" |
| ofm_layout = "NHCWB16" if get_expected and i != 2 else "NHWC" |
| x = infra.make_ethosu_conv2d( |
| ifm=x, |
| ifm_channels=8, |
| ofm_channels=8, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout=ifm_layout, |
| ofm_layout=ofm_layout, |
| ) |
| return relay.Function(relay.analysis.free_vars(x), x) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_multiple_depthwise_convolution(): |
| """Test layout optimization pass on multiple depthwise convolutions. |
| |
| depthwise_conv_1 |
| | |
| depthwise_conv_2 |
| | |
| depthwise_conv_3 |
| """ |
| |
| def get_graph(get_expected=False): |
| x = relay.var("x", shape=(1, 8, 8, 4), dtype="int8") |
| for i in range(3): |
| ifm_layout = "NHCWB16" if get_expected and i != 0 else "NHWC" |
| ofm_layout = "NHCWB16" if get_expected and i != 2 else "NHWC" |
| x = infra.make_ethosu_depthwise_conv2d( |
| ifm=x, |
| channels=4, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout=ifm_layout, |
| ofm_layout=ofm_layout, |
| ) |
| return relay.Function(relay.analysis.free_vars(x), x) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_ignore_transform_operations(): |
| """Test layout optimization pass ignores transform operations |
| such as reshape and strided slice. |
| |
| conv_1 |
| | |
| reshape |
| | |
| strided_slice |
| | |
| conv_2 |
| """ |
| |
| def get_graph(): |
| in_1 = relay.var("x", shape=(1, 16, 16, 8), dtype="int8") |
| conv_1 = infra.make_ethosu_conv2d( |
| ifm=in_1, |
| ifm_channels=8, |
| ofm_channels=8, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| reshape = relay.reshape(conv_1, (1, 16, 16, 8)) |
| strided_slice = relay.strided_slice(reshape, (0, 0, 0, 0), (1, 16, 16, 8)) |
| conv_2 = infra.make_ethosu_conv2d( |
| ifm=strided_slice, |
| ifm_channels=8, |
| ofm_channels=8, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(conv_2), conv_2) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_ignore_concatenate(): |
| """Test layout optimization pass ignores the concatenate operation, |
| when layout transformation cannot occur. |
| |
| in_1 in_2 |
| \ / |
| \ conv_1 |
| \ / |
| concat |
| | |
| conv_2 |
| """ |
| |
| def get_graph(): |
| in_1 = relay.var("x", shape=(1, 16, 16, 8), dtype="int8") |
| in_2 = relay.var("y", shape=(1, 16, 16, 8), dtype="int8") |
| conv_1 = infra.make_ethosu_conv2d( |
| ifm=in_2, |
| ifm_channels=8, |
| ofm_channels=8, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| concat = relay.concatenate([in_1, conv_1], axis=1) |
| conv_2 = infra.make_ethosu_conv2d( |
| ifm=concat, |
| ifm_channels=8, |
| ofm_channels=4, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(conv_2), conv_2) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_ignore_concatnate_with_layout_transform(): |
| """Test the layout optimization pass ignores the concatenate |
| operation and performs a layout transformation. |
| |
| in_1 in_2 |
| \ / |
| pool_1 pool_2 |
| \ / |
| concat |
| | |
| pool_3 |
| """ |
| |
| def get_graph(): |
| dtype = "int8" |
| |
| in_1 = relay.var("x", shape=(1, 16, 16, 8), dtype=dtype) |
| in_2 = relay.var("y", shape=(1, 16, 16, 8), dtype=dtype) |
| pool_1 = infra.make_ethosu_pooling( |
| in_1, |
| "MAX", |
| (1, 1), |
| ofm_channels=8, |
| ofm_dtype=dtype, |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| pool_2 = infra.make_ethosu_pooling( |
| in_2, |
| "MAX", |
| (1, 1), |
| ofm_channels=8, |
| ofm_dtype=dtype, |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| concat = relay.concatenate([pool_1, pool_2], axis=1) |
| pool_3 = infra.make_ethosu_pooling( |
| concat, |
| "MAX", |
| (1, 1), |
| ofm_channels=8, |
| ofm_dtype=dtype, |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(pool_3), pool_3) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_multiple_inputs(): |
| """Test the layout optimization pass works as expected when there |
| are multiple inputs in the graph. |
| |
| pool_1 pool_2 pool_3 |
| \ | / |
| \ | / |
| concat |
| | |
| conv |
| """ |
| |
| def get_graph(): |
| poolings = [] |
| for _ in range(3): |
| dtype = "int8" |
| |
| inp = relay.var("x", shape=(1, 3, 3, 4), dtype=dtype) |
| pool = infra.make_ethosu_pooling( |
| inp, |
| "MAX", |
| (1, 1), |
| ofm_channels=4, |
| ofm_dtype=dtype, |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| poolings.append(pool) |
| concat = relay.concatenate(poolings, axis=0) |
| conv = infra.make_ethosu_conv2d( |
| ifm=concat, |
| ifm_channels=8, |
| ofm_channels=4, |
| kernel_shape=(1, 1), |
| padding=(0, 0), |
| strides=(1, 1), |
| dilation=(1, 1), |
| ifm_layout="NHWC", |
| ofm_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(conv), conv) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_multiple_outputs(): |
| """Test the layout optimization pass works as expected when there |
| are multiple outputs in the graph. |
| |
| pool_1 |
| / | \ |
| pool_2 pool_3 pool_4 |
| \ | / |
| concat |
| """ |
| |
| def get_graph(get_expected=False): |
| dtype = "int8" |
| |
| in_1 = relay.var("x", shape=(1, 4, 4, 8), dtype=dtype) |
| pool_1 = infra.make_ethosu_pooling( |
| in_1, |
| "MAX", |
| (1, 1), |
| ofm_channels=4, |
| ofm_dtype=dtype, |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout="NHWC", |
| ofm_layout="NHCWB16" if get_expected else "NHWC", |
| ) |
| poolings = [] |
| for _ in range(3): |
| poolings.append( |
| infra.make_ethosu_pooling( |
| pool_1, |
| "MAX", |
| (1, 1), |
| ofm_channels=4, |
| ofm_dtype=dtype, |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout="NHCWB16" if get_expected else "NHWC", |
| ofm_layout="NHWC", |
| ) |
| ) |
| concat = relay.concatenate(poolings, axis=0) |
| return relay.Function(relay.analysis.free_vars(concat), concat) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_multiple_binary_elementwise(): |
| """Test the layout optimization pass works as expected for |
| binary elementwise operations. |
| |
| add_1 add_2 |
| \ / |
| \ / |
| add_3 |
| """ |
| |
| def get_graph(get_expected=False): |
| in_1 = relay.var("x", shape=(1, 2, 2, 2), dtype="int8") |
| in_2 = relay.var("y", shape=(1, 2, 2, 2), dtype="int8") |
| in_3 = relay.var("z", shape=(1, 2, 2, 2), dtype="int8") |
| add_1 = infra.make_ethosu_binary_elementwise( |
| in_1, |
| in_2, |
| ifm_channels=2, |
| ifm2_channels=2, |
| operator_type="ADD", |
| ofm_dtype="int8", |
| ifm_layout="NHWC", |
| ifm2_layout="NHWC", |
| ofm_layout="NHCWB16" if get_expected else "NHWC", |
| ) |
| add_2 = infra.make_ethosu_binary_elementwise( |
| in_2, |
| in_3, |
| ifm_channels=2, |
| ifm2_channels=2, |
| operator_type="ADD", |
| ofm_dtype="int8", |
| ifm_layout="NHWC", |
| ifm2_layout="NHWC", |
| ofm_layout="NHCWB16" if get_expected else "NHWC", |
| ) |
| add_3 = infra.make_ethosu_binary_elementwise( |
| add_1, |
| add_2, |
| ifm_channels=2, |
| ifm2_channels=2, |
| operator_type="ADD", |
| ofm_dtype="int8", |
| ifm_layout="NHCWB16" if get_expected else "NHWC", |
| ifm2_layout="NHCWB16" if get_expected else "NHWC", |
| ofm_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(add_3), add_3) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_multiple_pooling(): |
| """Test the layout optimization pass works as expected for |
| multiple pooling operations. |
| |
| pool_1 |
| | |
| pool_2 |
| | |
| pool_3 |
| """ |
| |
| def get_graph(get_expected=False): |
| dtype = "int8" |
| |
| x = relay.var("x", shape=(1, 8, 8, 4), dtype=dtype) |
| for i in range(3): |
| ifm_layout = "NHCWB16" if get_expected and i != 0 else "NHWC" |
| ofm_layout = "NHCWB16" if get_expected and i != 2 else "NHWC" |
| x = infra.make_ethosu_pooling( |
| x, |
| "MAX", |
| (1, 1), |
| ofm_channels=4, |
| ofm_dtype=dtype, |
| strides=(1, 1), |
| padding=(0, 0), |
| ifm_layout=ifm_layout, |
| ofm_layout=ofm_layout, |
| ) |
| return relay.Function(relay.analysis.free_vars(x), x) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_multiple_unary_elementwise(): |
| """Test the layout optimization pass works as expected for multiple |
| unary elementwise operations. |
| |
| abs_1 |
| | |
| abs_2 |
| | |
| abs_3 |
| """ |
| |
| def get_graph(get_expected=False): |
| x = relay.var("x", shape=(1, 8, 8, 4), dtype="int8") |
| for i in range(3): |
| ifm_layout = "NHCWB16" if get_expected and i != 0 else "NHWC" |
| ofm_layout = "NHCWB16" if get_expected and i != 2 else "NHWC" |
| x = infra.make_ethosu_unary_elementwise( |
| x, |
| ofm_channels=4, |
| operator_type="ABS", |
| ifm_layout=ifm_layout, |
| ofm_layout=ofm_layout, |
| ) |
| return relay.Function(relay.analysis.free_vars(x), x) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_op_without_ethosu_consumer(): |
| """Test the layout optimization pass works as expected when |
| there is a case that the output layout should not be altered |
| since not all consumers are NPU operations (in this case conv). |
| |
| depthwise |
| | |
| conv |
| / \ |
| | pool |
| \ / |
| (concat) |
| """ |
| |
| def get_graph(get_expected=False): |
| exp_layout = "NHCWB16" if get_expected else "NHWC" |
| dtype = "int8" |
| |
| x = relay.var("x", shape=(1, 2, 2, 2), dtype=dtype) |
| depthwise = infra.make_ethosu_depthwise_conv2d( |
| x, 2, (1, 1), (0, 0), (1, 1), (0, 0), ofm_layout=exp_layout |
| ) |
| conv = infra.make_ethosu_conv2d( |
| depthwise, |
| 2, |
| 2, |
| (1, 1), |
| (0, 0), |
| (1, 1), |
| (0, 0), |
| ifm_layout=exp_layout, |
| ) |
| pool = infra.make_ethosu_pooling(conv, "MAX", (1, 1), 2, dtype, (1, 1), (0, 0)) |
| concat = relay.concatenate([conv, pool], axis=0) |
| return relay.Function(relay.analysis.free_vars(concat), concat) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_diamond_graph(): |
| """ |
| Test the layout optimizer pass works as expected on a diamond graph |
| with a case where the operation dominating the output operation |
| cannot be altered, but operations within the diamond can. |
| |
| pool_1 |
| | |
| pool_2 |
| / \ |
| | pool_3 |
| | | |
| | pool_4 |
| | | |
| | pool_5 |
| \ / |
| (concat) |
| """ |
| |
| def get_graph(get_expected=False): |
| exp_layout = "NHCWB16" if get_expected else "NHWC" |
| dtype = "int8" |
| |
| x = relay.var("x", shape=(1, 2, 2, 2), dtype=dtype) |
| pool_1 = infra.make_ethosu_pooling( |
| x, "MAX", (1, 1), 2, dtype, (1, 1), (0, 0), ofm_layout=exp_layout |
| ) |
| pool_2 = infra.make_ethosu_pooling( |
| pool_1, "MAX", (1, 1), 2, dtype, (1, 1), (0, 0), ifm_layout=exp_layout |
| ) |
| pool_3 = infra.make_ethosu_pooling( |
| pool_2, "MAX", (1, 1), 2, dtype, (1, 1), (0, 0), ofm_layout=exp_layout |
| ) |
| pool_4 = infra.make_ethosu_pooling( |
| pool_3, |
| "MAX", |
| (1, 1), |
| 2, |
| dtype, |
| (1, 1), |
| (0, 0), |
| ifm_layout=exp_layout, |
| ofm_layout=exp_layout, |
| ) |
| pool_5 = infra.make_ethosu_pooling( |
| pool_4, "MAX", (1, 1), 2, dtype, (1, 1), (0, 0), ifm_layout=exp_layout |
| ) |
| concat = relay.concatenate([pool_2, pool_5], axis=0) |
| return relay.Function(relay.analysis.free_vars(concat), concat) |
| |
| a = _optimize(get_graph()) |
| b = _optimize(get_graph(get_expected=True), optimize=False) |
| _assert_structural_equal(a, b) |
| |
| |
| def test_same_output_multiple_convolutions(): |
| """Test running the layout optimization pass with multiple convolutions |
| gives same output as TFLite.""" |
| |
| np.random.seed(0) |
| dtype = "int8" |
| ifm_shape = (1, 8, 8, 32) |
| kernel_shape = (1, 1, 32, 32) |
| |
| def create_model(): |
| class Model(tf.Module): |
| @tf.function |
| def tf_function(self, x): |
| for _ in range(3): |
| x = tf.nn.conv2d( |
| x, |
| filters=tf.constant(np.random.uniform(size=kernel_shape), dtype=tf.float32), |
| strides=(1, 1), |
| padding="SAME", |
| data_format="NHWC", |
| dilations=1, |
| ) |
| return x |
| |
| model = Model() |
| concrete_func = model.tf_function.get_concrete_function( |
| tf.TensorSpec(ifm_shape, dtype=tf.float32) |
| ) |
| |
| # Convert the model |
| def representative_dataset(): |
| for _ in range(100): |
| data = np.random.rand(*tuple(ifm_shape)) |
| yield [data.astype(np.float32)] |
| |
| converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func]) |
| converter.optimizations = [tf.lite.Optimize.DEFAULT] |
| converter.representative_dataset = representative_dataset |
| converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] |
| converter.inference_input_type = tf.int8 |
| converter.inference_output_type = tf.int8 |
| return converter.convert() |
| |
| _compile_and_compare_model(create_model(), ifm_shape, dtype) |
| |
| |
| def test_same_output_multiple_pooling(): |
| """Test running the layout optimization pass with multiple pooling |
| operations gives same output as TFLite.""" |
| |
| np.random.seed(0) |
| dtype = "int8" |
| ifm_shape = (1, 4, 2, 7) |
| |
| def create_model(): |
| class Model(tf.Module): |
| @tf.function |
| def tf_function(self, x): |
| for _ in range(2): |
| x = tf.nn.max_pool2d(x, (1, 1), (1, 1), "SAME", "NHWC") |
| return x |
| |
| model = Model() |
| concrete_func = model.tf_function.get_concrete_function( |
| tf.TensorSpec(ifm_shape, dtype=tf.float32) |
| ) |
| |
| # Convert the model |
| def representative_dataset(): |
| for _ in range(100): |
| data = np.random.rand(*tuple(ifm_shape)) |
| yield [data.astype(np.float32)] |
| |
| converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func]) |
| converter.optimizations = [tf.lite.Optimize.DEFAULT] |
| converter.representative_dataset = representative_dataset |
| converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] |
| converter.inference_input_type = tf.int8 |
| converter.inference_output_type = tf.int8 |
| return converter.convert() |
| |
| _compile_and_compare_model(create_model(), ifm_shape, dtype) |
| |
| |
| def test_layout_optimizer_runs_in_compilation_pipeline(): |
| """Checks that the layout optimization pass runs as part of the NPU compilation |
| pipeline.""" |
| |
| def get_graph(): |
| x = relay.var("x", shape=(1, 4, 4, 4), dtype="int8") |
| for _ in range(2): |
| x = relay.nn.max_pool2d(x, layout="NHWC") |
| |
| func = relay.Function(relay.analysis.free_vars(x), x) |
| return tvm.IRModule.from_expr(func) |
| |
| mod = get_graph() |
| mod = partition_for_ethosu(mod) |
| mod = relay_to_tir(mod) |
| |
| external_gv_name = mod["main"].body.op.name_hint |
| prim_func = mod[external_gv_name] |
| |
| # Check for hints in the TIR prim func that the layout optimization pass has ran |
| ops = prim_func.body.body.seq |
| max_pool1, max_pool2 = ops |
| |
| assert str(max_pool1.value.args[31]) == '"NHCWB16"' |
| assert str(max_pool2.value.args[14]) == '"NHCWB16"' |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
