blob: 7fdc4bd8843a59ae8b0df873125353793944f08d [file] [log] [blame]
Diqing Zhong94457b12020-12-09 15:22:40 +01001# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17# Description:
18# Unit tests for graph_optimiser
19import numpy as np
20
21from ethosu.vela.graph_optimiser import convert_batched_fc_shape
22from ethosu.vela.operation import Op
23from ethosu.vela.tensor import create_const_tensor
patrik.gustavssoneeb85152020-12-21 17:10:40 +000024from ethosu.vela.tensor import Shape4D
Diqing Zhong94457b12020-12-09 15:22:40 +010025from ethosu.vela.tensor import Tensor
26from ethosu.vela.test import testutil
27
28
29def test_convert_batched_fc():
30 """Tests shape conversion of batched fully connected"""
31 shape = [4, 8]
32 ifm = create_const_tensor("test_in", shape, np.uint8, np.zeros(shape))
33 weights = create_const_tensor("weight_in", shape, np.uint8, np.zeros(shape))
34 ofm = Tensor(ifm.shape, np.uint8, "test_out")
35 op = testutil.create_op(Op.FullyConnected, [ifm, weights], ofm)
Patrik Gustavsson2349d422020-12-01 16:02:29 +010036
Diqing Zhong94457b12020-12-09 15:22:40 +010037 ifm.consumer_list.append(op)
38
patrik.gustavssoneeb85152020-12-21 17:10:40 +000039 op.ifm_shapes.append(Shape4D([4, 1, 1, 8]))
40 op.ofm_shapes.append(Shape4D([4, 1, 1, 8]))
Patrik Gustavsson2349d422020-12-01 16:02:29 +010041
Diqing Zhong94457b12020-12-09 15:22:40 +010042 prev_op = op.clone()
Patrik Gustavsson2349d422020-12-01 16:02:29 +010043 prev_op.ifm_shapes = op.ifm_shapes
44 prev_op.ofm_shapes = op.ofm_shapes
45
Diqing Zhong94457b12020-12-09 15:22:40 +010046 conv_op = convert_batched_fc_shape(op, None, None)
47
48 assert conv_op.ifm != prev_op.ifm
49 assert conv_op.ofm != prev_op.ofm
50 assert conv_op.type == Op.FullyConnected
51 assert len(conv_op.ifm.shape) == 4
52 assert conv_op.ifm.shape == conv_op.ofm.shape
53 assert conv_op.ifm.ops[0].type == Op.Reshape
54
55 shape = [1, 8]
56 ifm.shape = shape
57 weights.shape = shape
58 ofm.shape = shape
59 op = testutil.create_op(Op.FullyConnected, [ifm, weights], ofm)
60 ifm.consumer_list.append(op)
61
Patrik Gustavsson2349d422020-12-01 16:02:29 +010062 op.ifm_shapes.append([1, 1, 1, 8])
63 op.ofm_shapes.append([1, 1, 1, 8])
64
Diqing Zhong94457b12020-12-09 15:22:40 +010065 prev_op = op.clone()
Patrik Gustavsson2349d422020-12-01 16:02:29 +010066 prev_op.ifm_shapes = op.ifm_shapes
67 prev_op.ofm_shapes = op.ofm_shapes
68
Diqing Zhong94457b12020-12-09 15:22:40 +010069 conv_op = convert_batched_fc_shape(op, None, None)
70
71 assert conv_op.ifm == prev_op.ifm
72 assert conv_op.ofm == prev_op.ofm
73 assert conv_op.type == Op.FullyConnected
74 assert len(conv_op.ifm.shape) == 2
75 assert conv_op.ifm.shape == conv_op.ofm.shape