blob: 2fc7622ca3f26a6dc50a6e62c0d30a8a73baaa64 [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
19
20from .data_type import DataType
21from .high_level_command_to_npu_op import ifm_ifm2_correct_order
22from .operation import ActivationFunction
23from .operation import Op
24from .operation import Operation
25from .tensor import create_reshape_tensor
26from .tensor import QuantizationParameters
27from .tensor import Tensor
28
29
30def create_avgpool_nop(name: str) -> Operation:
31 op = Operation(Op.AvgPool, name)
32 op.attrs["padding"] = b"VALID"
33 op.attrs["stride_w"] = 1
34 op.attrs["stride_h"] = 1
35 op.attrs["filter_width"] = 1
36 op.attrs["filter_height"] = 1
37 op.attrs["strides"] = [1, 1, 1, 1]
38 op.attrs["ksize"] = [1, 1, 1, 1]
39 op.attrs["skirt"] = [0, 0, 0, 0]
40 op.attrs["explicit_padding"] = [0, 0, 0, 0]
41 return op
42
43
44def create_depthwise_maxpool(
45 name: str, ifm: Tensor, quantization: QuantizationParameters, activation: Optional[ActivationFunction] = None
46) -> Operation:
47 op = Operation(Op.MaxPool, name)
48 height = ifm.shape[1] * ifm.shape[2]
49 width = ifm.shape[3]
50 ifm_shape = [1, height, width, 1]
51 op.attrs["padding"] = b"VALID"
52 op.attrs["stride_w"] = 1
53 op.attrs["stride_h"] = 1
54 op.attrs["filter_width"] = width
55 op.attrs["filter_height"] = 1
56 op.attrs["strides"] = [1, op.attrs["stride_h"], op.attrs["stride_w"], 1]
57 op.attrs["ksize"] = [1, op.attrs["filter_height"], op.attrs["filter_width"], 1]
58 op.activation = activation
59 op.inputs = [create_reshape_tensor(ifm, ifm_shape)]
60 ofm = Tensor([1, height, 1, 1], ifm.dtype, op.name + "_tens0")
61 ofm.quantization = quantization
62 op.set_output_tensor(ofm)
63 return op
64
65
66def create_reduce_sum(
67 name: str, ifm: Tensor, quantization: QuantizationParameters, activation: Optional[ActivationFunction] = None
68) -> Operation:
69 op = Operation(Op.ReduceSum, name)
70 op.attrs["padding"] = b"VALID"
71 op.attrs["stride_w"] = 1
72 op.attrs["stride_h"] = 1
73 op.attrs["filter_width"] = 1
74 op.attrs["filter_height"] = 1
75 op.attrs["strides"] = [1, op.attrs["stride_h"], op.attrs["stride_w"], 1]
76 op.attrs["ksize"] = [1, op.attrs["filter_height"], op.attrs["filter_width"], 1]
77 op.add_input_tensor(ifm)
78 op.activation = activation
79 ofm_shape = [1, ifm.shape[1], ifm.shape[2], 1]
80 sum_of_exp = Tensor(ofm_shape, DataType.int32, op.name + "_tens0")
81 sum_of_exp.quantization = quantization
82 op.set_output_tensor(sum_of_exp)
83 return op
84
85
86def create_add(
87 name: str,
88 ifm: Tensor,
89 ifm2: Tensor,
90 quantization: QuantizationParameters,
91 activation: Optional[ActivationFunction] = None,
92 dtype: Optional[DataType] = None,
93 attrs: Optional[dict] = None,
94) -> Operation:
95 return create_binary_elementwise(Op.Add, name, ifm, ifm2, quantization, activation, dtype, attrs)
96
97
98def create_clz(
99 name: str,
100 ifm: Tensor,
101 quantization: QuantizationParameters,
102 activation: Optional[ActivationFunction] = None,
103 dtype: Optional[DataType] = None,
104 attrs: Optional[dict] = None,
105) -> Operation:
106 return create_unary_elementwise(Op.CLZ, name, ifm, quantization, activation, dtype, attrs)
107
108
109def create_mul(
110 name: str,
111 ifm: Tensor,
112 ifm2: Tensor,
113 quantization: QuantizationParameters,
114 activation: Optional[ActivationFunction] = None,
115 dtype: Optional[DataType] = None,
116 attrs: Optional[dict] = None,
117) -> Operation:
118 return create_binary_elementwise(Op.Mul, name, ifm, ifm2, quantization, activation, dtype, attrs)
119
120
121def create_shl(
122 name: str,
123 ifm: Tensor,
124 ifm2: Tensor,
125 quantization: QuantizationParameters,
126 activation: Optional[ActivationFunction] = None,
127 dtype: Optional[DataType] = None,
128 attrs: Optional[dict] = None,
129) -> Operation:
130 return create_binary_elementwise(Op.SHL, name, ifm, ifm2, quantization, activation, dtype, attrs)
131
132
133def create_shr(
134 name: str,
135 ifm: Tensor,
136 ifm2: Tensor,
137 quantization: QuantizationParameters,
138 activation: Optional[ActivationFunction] = None,
139 dtype: Optional[DataType] = None,
140 attrs: Optional[dict] = None,
141) -> Operation:
142 return create_binary_elementwise(Op.SHR, name, ifm, ifm2, quantization, activation, dtype, attrs)
143
144
145def create_sub(
146 name: str,
147 ifm: Tensor,
148 ifm2: Tensor,
149 quantization: QuantizationParameters,
150 activation: Optional[ActivationFunction] = None,
151 dtype: Optional[DataType] = None,
152 attrs: Optional[dict] = None,
153) -> Operation:
154 return create_binary_elementwise(Op.Sub, name, ifm, ifm2, quantization, activation, dtype, attrs)
155
156
157def create_unary_elementwise(
158 op_type: Op,
159 name: str,
160 ifm: Tensor,
161 quantization: QuantizationParameters,
162 activation: Optional[ActivationFunction] = None,
163 dtype: Optional[DataType] = None,
164 attrs: Optional[dict] = None,
165) -> Operation:
166 return create_binary_elementwise(op_type, name, ifm, None, quantization, activation, dtype, attrs)
167
168
169def create_binary_elementwise(
170 op_type: Op,
171 name: str,
172 ifm: Tensor,
173 ifm2: Tensor,
174 quantization: QuantizationParameters,
175 activation: Optional[ActivationFunction] = None,
176 dtype: Optional[DataType] = None,
177 attrs: Optional[dict] = None,
178) -> Operation:
179 op = Operation(op_type, name)
180 op.add_input_tensor(ifm)
181 if ifm2:
182 op.add_input_tensor(ifm2)
183 op.activation = activation
184 if not dtype:
185 dtype = ifm.dtype
186 if attrs:
187 op.attrs.update(attrs)
188 ofm_shape = ifm.shape if ifm2 is None or ifm_ifm2_correct_order(ifm.shape, ifm2.shape) else ifm2.shape
189 ofm = Tensor(ofm_shape, dtype, f"{op.name}_tens0")
190 ofm.quantization = quantization
191 op.set_output_tensor(ofm)
192 return op