blob: 7015b7993d3b25ac63e0ecb2d50d6d11a4a2a8c9 [file] [log] [blame]
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +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# Description:
17# Utility functions for creating Network Operations.
18from typing import Optional
Fredrik Svedberge82be7c2021-01-18 15:21:03 +010019from typing import Tuple
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010020
21from .data_type import DataType
22from .high_level_command_to_npu_op import ifm_ifm2_correct_order
23from .operation import ActivationFunction
24from .operation import Op
25from .operation import Operation
Michael McGeagh16895482020-12-14 15:51:20 +000026from .operation import Padding
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010027from .tensor import create_reshape_tensor
28from .tensor import QuantizationParameters
29from .tensor import Tensor
30
31
32def create_avgpool_nop(name: str) -> Operation:
33 op = Operation(Op.AvgPool, name)
Michael McGeagh16895482020-12-14 15:51:20 +000034 op.attrs["padding"] = Padding.VALID
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010035 op.attrs["stride_w"] = 1
36 op.attrs["stride_h"] = 1
37 op.attrs["filter_width"] = 1
38 op.attrs["filter_height"] = 1
39 op.attrs["strides"] = [1, 1, 1, 1]
40 op.attrs["ksize"] = [1, 1, 1, 1]
41 op.attrs["skirt"] = [0, 0, 0, 0]
42 op.attrs["explicit_padding"] = [0, 0, 0, 0]
43 return op
44
45
46def create_depthwise_maxpool(
47 name: str, ifm: Tensor, quantization: QuantizationParameters, activation: Optional[ActivationFunction] = None
48) -> Operation:
49 op = Operation(Op.MaxPool, name)
50 height = ifm.shape[1] * ifm.shape[2]
51 width = ifm.shape[3]
52 ifm_shape = [1, height, width, 1]
Michael McGeagh16895482020-12-14 15:51:20 +000053 op.attrs["padding"] = Padding.VALID
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010054 op.attrs["stride_w"] = 1
55 op.attrs["stride_h"] = 1
56 op.attrs["filter_width"] = width
57 op.attrs["filter_height"] = 1
58 op.attrs["strides"] = [1, op.attrs["stride_h"], op.attrs["stride_w"], 1]
59 op.attrs["ksize"] = [1, op.attrs["filter_height"], op.attrs["filter_width"], 1]
60 op.activation = activation
61 op.inputs = [create_reshape_tensor(ifm, ifm_shape)]
62 ofm = Tensor([1, height, 1, 1], ifm.dtype, op.name + "_tens0")
63 ofm.quantization = quantization
64 op.set_output_tensor(ofm)
Patrik Gustavsson2349d422020-12-01 16:02:29 +010065 op.set_ifm_ofm_shapes()
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010066 return op
67
68
69def create_reduce_sum(
70 name: str, ifm: Tensor, quantization: QuantizationParameters, activation: Optional[ActivationFunction] = None
71) -> Operation:
72 op = Operation(Op.ReduceSum, name)
Michael McGeagh16895482020-12-14 15:51:20 +000073 op.attrs["padding"] = Padding.VALID
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010074 op.attrs["stride_w"] = 1
75 op.attrs["stride_h"] = 1
76 op.attrs["filter_width"] = 1
77 op.attrs["filter_height"] = 1
78 op.attrs["strides"] = [1, op.attrs["stride_h"], op.attrs["stride_w"], 1]
79 op.attrs["ksize"] = [1, op.attrs["filter_height"], op.attrs["filter_width"], 1]
80 op.add_input_tensor(ifm)
81 op.activation = activation
82 ofm_shape = [1, ifm.shape[1], ifm.shape[2], 1]
83 sum_of_exp = Tensor(ofm_shape, DataType.int32, op.name + "_tens0")
84 sum_of_exp.quantization = quantization
85 op.set_output_tensor(sum_of_exp)
Patrik Gustavsson2349d422020-12-01 16:02:29 +010086 op.set_ifm_ofm_shapes()
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010087 return op
88
89
90def create_add(
91 name: str,
92 ifm: Tensor,
93 ifm2: Tensor,
94 quantization: QuantizationParameters,
95 activation: Optional[ActivationFunction] = None,
96 dtype: Optional[DataType] = None,
97 attrs: Optional[dict] = None,
98) -> Operation:
99 return create_binary_elementwise(Op.Add, name, ifm, ifm2, quantization, activation, dtype, attrs)
100
101
Fredrik Svedberge82be7c2021-01-18 15:21:03 +0100102def create_rescale_add(
103 name: str,
104 ifm: Tensor,
105 ifm2: Tensor,
106 rescale: Tuple[int, int],
107 quantization: QuantizationParameters,
108 activation: Optional[ActivationFunction] = None,
109 dtype: Optional[DataType] = None,
110 attrs: Optional[dict] = None,
111) -> Operation:
112 op = create_binary_elementwise(Op.RescaleAdd, name, ifm, ifm2, quantization, activation, dtype, attrs)
113 op.rescale = rescale
114 return op
115
116
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +0100117def create_clz(
118 name: str,
119 ifm: Tensor,
120 quantization: QuantizationParameters,
121 activation: Optional[ActivationFunction] = None,
122 dtype: Optional[DataType] = None,
123 attrs: Optional[dict] = None,
124) -> Operation:
125 return create_unary_elementwise(Op.CLZ, name, ifm, quantization, activation, dtype, attrs)
126
127
128def create_mul(
129 name: str,
130 ifm: Tensor,
131 ifm2: Tensor,
132 quantization: QuantizationParameters,
133 activation: Optional[ActivationFunction] = None,
134 dtype: Optional[DataType] = None,
135 attrs: Optional[dict] = None,
136) -> Operation:
137 return create_binary_elementwise(Op.Mul, name, ifm, ifm2, quantization, activation, dtype, attrs)
138
139
140def create_shl(
141 name: str,
142 ifm: Tensor,
143 ifm2: Tensor,
144 quantization: QuantizationParameters,
145 activation: Optional[ActivationFunction] = None,
146 dtype: Optional[DataType] = None,
147 attrs: Optional[dict] = None,
148) -> Operation:
149 return create_binary_elementwise(Op.SHL, name, ifm, ifm2, quantization, activation, dtype, attrs)
150
151
152def create_shr(
153 name: str,
154 ifm: Tensor,
155 ifm2: Tensor,
156 quantization: QuantizationParameters,
157 activation: Optional[ActivationFunction] = None,
158 dtype: Optional[DataType] = None,
159 attrs: Optional[dict] = None,
160) -> Operation:
161 return create_binary_elementwise(Op.SHR, name, ifm, ifm2, quantization, activation, dtype, attrs)
162
163
164def create_sub(
165 name: str,
166 ifm: Tensor,
167 ifm2: Tensor,
168 quantization: QuantizationParameters,
169 activation: Optional[ActivationFunction] = None,
170 dtype: Optional[DataType] = None,
171 attrs: Optional[dict] = None,
172) -> Operation:
173 return create_binary_elementwise(Op.Sub, name, ifm, ifm2, quantization, activation, dtype, attrs)
174
175
176def create_unary_elementwise(
177 op_type: Op,
178 name: str,
179 ifm: Tensor,
180 quantization: QuantizationParameters,
181 activation: Optional[ActivationFunction] = None,
182 dtype: Optional[DataType] = None,
183 attrs: Optional[dict] = None,
184) -> Operation:
185 return create_binary_elementwise(op_type, name, ifm, None, quantization, activation, dtype, attrs)
186
187
188def create_binary_elementwise(
189 op_type: Op,
190 name: str,
191 ifm: Tensor,
192 ifm2: Tensor,
193 quantization: QuantizationParameters,
194 activation: Optional[ActivationFunction] = None,
195 dtype: Optional[DataType] = None,
196 attrs: Optional[dict] = None,
197) -> Operation:
198 op = Operation(op_type, name)
199 op.add_input_tensor(ifm)
200 if ifm2:
201 op.add_input_tensor(ifm2)
202 op.activation = activation
203 if not dtype:
204 dtype = ifm.dtype
205 if attrs:
206 op.attrs.update(attrs)
207 ofm_shape = ifm.shape if ifm2 is None or ifm_ifm2_correct_order(ifm.shape, ifm2.shape) else ifm2.shape
208 ofm = Tensor(ofm_shape, dtype, f"{op.name}_tens0")
209 ofm.quantization = quantization
210 op.set_output_tensor(ofm)
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100211 op.set_ifm_ofm_shapes()
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +0100212 return op