|
| 1 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +# or more contributor license agreements. See the NOTICE file |
| 3 | +# distributed with this work for additional information |
| 4 | +# regarding copyright ownership. The ASF licenses this file |
| 5 | +# to you under the Apache License, Version 2.0 (the |
| 6 | +# "License"); you may not use this file except in compliance |
| 7 | +# with the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | +# pylint: disable=wrong-import-position,invalid-name |
| 18 | + |
| 19 | +""" |
| 20 | +Test the cascader in the compilation flow. |
| 21 | +""" |
| 22 | + |
| 23 | +import pytest |
| 24 | + |
| 25 | +pytest.importorskip("ethosu.vela") |
| 26 | + |
| 27 | +import numpy as np |
| 28 | + |
| 29 | +import tvm |
| 30 | +from tvm import relay |
| 31 | +from tvm.relay.backend.contrib.ethosu.codegen import _create_cascader |
| 32 | +from tvm.relay.backend.contrib.ethosu.tir.compiler import _lower_to_tir |
| 33 | +from tvm.contrib.ethosu.cascader import MemoryRegion, EthosuDeviceConfig |
| 34 | + |
| 35 | +from .. import infra as test_infra |
| 36 | +from . import infra as cascader_test_infra |
| 37 | + |
| 38 | + |
| 39 | +def _ethos_u55_cascader(): |
| 40 | + sram = MemoryRegion( |
| 41 | + name="SRAM", |
| 42 | + size=10**6, |
| 43 | + read_bandwidth=16, |
| 44 | + write_bandwidth=16, |
| 45 | + read_latency=0, |
| 46 | + write_latency=0, |
| 47 | + burst_length=1, |
| 48 | + ) |
| 49 | + flash = MemoryRegion(name="FLASH", size=10**7, read_bandwidth=4, write_bandwidth=4) |
| 50 | + |
| 51 | + device_config = EthosuDeviceConfig("ethos-u55-256") |
| 52 | + cascader_options = cascader_test_infra.make_options( |
| 53 | + cascade_region=sram, |
| 54 | + max_proposals=64, |
| 55 | + stripe_factors=4, |
| 56 | + max_plan_size=10, |
| 57 | + max_open_plans=8, |
| 58 | + max_closed_plans=32, |
| 59 | + always_copy_size=1024, |
| 60 | + disable_pareto_plans=False, |
| 61 | + disable_pareto_proposals=False, |
| 62 | + enable_striping=False, |
| 63 | + ) |
| 64 | + return _create_cascader( |
| 65 | + options=cascader_options, |
| 66 | + io_region=sram, |
| 67 | + constant_region=flash, |
| 68 | + working_regions=[sram], |
| 69 | + device_config=device_config, |
| 70 | + ) |
| 71 | + |
| 72 | + |
| 73 | +def _compile_model(relay_function): |
| 74 | + mod = tvm.IRModule() |
| 75 | + mod["main"] = relay_function |
| 76 | + mod = relay.transform.InferType()(mod) |
| 77 | + tir_mod = _lower_to_tir(mod["main"], _ethos_u55_cascader())[0] |
| 78 | + return tir_mod["main"] |
| 79 | + |
| 80 | + |
| 81 | +def _create_single_conv2d(): |
| 82 | + ifm = relay.var("x", shape=(1, 8, 8, 4), dtype="int8") |
| 83 | + conv1 = test_infra.make_ethosu_conv2d(ifm, 4, 4, (3, 3), (1, 1), (1, 1), (1, 1)) |
| 84 | + func = relay.Function(relay.analysis.free_vars(conv1), conv1) |
| 85 | + return func |
| 86 | + |
| 87 | + |
| 88 | +def _create_double_conv2d(): |
| 89 | + ifm = relay.var("x", shape=(1, 8, 8, 4), dtype="int8") |
| 90 | + conv1 = test_infra.make_ethosu_conv2d(ifm, 4, 4, (3, 3), (1, 1), (1, 1), (1, 1)) |
| 91 | + conv2 = test_infra.make_ethosu_conv2d(conv1, 4, 4, (1, 3), (1, 1), (1, 1), (1, 1)) |
| 92 | + func = relay.Function(relay.analysis.free_vars(conv2), conv2) |
| 93 | + return func |
| 94 | + |
| 95 | + |
| 96 | +def _create_scalar_add(): |
| 97 | + ifm = relay.var("x", shape=(1, 5, 4, 3), dtype="int8") |
| 98 | + ifm2 = relay.const(np.ones((1, 1, 1, 1)), dtype="int8") |
| 99 | + add = test_infra.make_ethosu_binary_elementwise( |
| 100 | + ifm, ifm2, ifm_channels=3, ifm2_channels=1, operator_type="ADD", ofm_dtype="int8" |
| 101 | + ) |
| 102 | + func = relay.Function(relay.analysis.free_vars(add), add) |
| 103 | + return func |
| 104 | + |
| 105 | + |
| 106 | +def test_single_conv_compute_cycles_hint(): |
| 107 | + """ |
| 108 | + Check the "compute_cycles_hint" annotation remains in the lowering flow |
| 109 | + for single convolution. |
| 110 | + """ |
| 111 | + primfunc = _compile_model(_create_single_conv2d()) |
| 112 | + ops = primfunc.body.body.body.seq |
| 113 | + |
| 114 | + compute_cycles_hints = [2304, 640, 320] |
| 115 | + for op, compute_cycle_hint in zip(ops, compute_cycles_hints): |
| 116 | + assert op.attr_key == "pragma_compute_cycles_hint" |
| 117 | + assert op.value == compute_cycle_hint |
| 118 | + |
| 119 | + |
| 120 | +def test_double_conv_compute_cycles_hint(): |
| 121 | + """ |
| 122 | + Check the "compute_cycles_hint" annotation remains in the lowering flow |
| 123 | + for double convolution. |
| 124 | + """ |
| 125 | + primfunc = _compile_model(_create_double_conv2d()) |
| 126 | + ops = primfunc.body.body.body.body.body.body.seq |
| 127 | + |
| 128 | + compute_cycles_hints = [2304, 640, 768, 640, 320, 240] |
| 129 | + for op, compute_cycle_hint in zip(ops, compute_cycles_hints): |
| 130 | + assert op.attr_key == "pragma_compute_cycles_hint" |
| 131 | + assert op.value == compute_cycle_hint |
| 132 | + |
| 133 | + |
| 134 | +def test_scalar_add_compute_cycles_hint(): |
| 135 | + """ |
| 136 | + Check the "compute_cycles_hint" annotation remains in the lowering flow |
| 137 | + for add with scalar values. |
| 138 | + """ |
| 139 | + primfunc = _compile_model(_create_scalar_add()) |
| 140 | + ops = primfunc.body.body.seq |
| 141 | + |
| 142 | + compute_cycles_hints = [16, 24] |
| 143 | + for op, compute_cycle_hint in zip(ops, compute_cycles_hints): |
| 144 | + assert op.attr_key == "pragma_compute_cycles_hint" |
| 145 | + assert op.value == compute_cycle_hint |
0 commit comments