blob: 00edf8356acdb4b8b6ec2b136af26efc703d6ab1 [file] [log] [blame]
Tim Hall79d07d22020-04-27 18:20:16 +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.
Tim Hall79d07d22020-04-27 18:20:16 +010016# Description:
17# Early optimisation of the network graph, using the rewrite_graph module to do the traversal of the graph. These are
18# split into two parts optimise_graph_a and optimise_graph_b.
Tim Hall79d07d22020-04-27 18:20:16 +010019import math
Diqing Zhong016b8272020-12-16 16:46:06 +010020import uuid
Diego Russoea6111a2020-04-14 18:41:58 +010021
22import numpy as np
23
Louis Verhaardd7911c42020-08-25 13:36:41 +020024from . import fp_math
Louis Verhaardb9fc33c2020-08-13 11:47:36 +020025from . import lut
Diego Russoea6111a2020-04-14 18:41:58 +010026from . import rewrite_graph
Louis Verhaardd7911c42020-08-25 13:36:41 +020027from . import scaling
Diego Russoea6111a2020-04-14 18:41:58 +010028from .data_type import DataType
Tim Halle6ccd872020-11-09 16:46:37 +000029from .debug_database import DebugDatabase
Louis Verhaard7db78962020-05-25 15:05:26 +020030from .errors import UnsupportedFeatureError
Dwight Lidman42fed942020-05-29 09:37:03 +020031from .ethos_u55_regs.ethos_u55_regs import resampling_mode
Louis Verhaard8912c532020-09-30 12:11:49 +020032from .numeric_util import clamp_sigmoid
Louis Verhaarde0ef2732020-06-03 08:56:44 +020033from .numeric_util import full_shape
Louis Verhaardf03bad32020-09-25 08:30:44 +020034from .numeric_util import round_away_zero
Louis Verhaarde8a5a782020-11-02 18:04:27 +010035from .operation import create_activation_function
Diego Russoe8a10452020-04-21 17:39:10 +010036from .operation import NpuBlockType
Louis Verhaardaee5d752020-09-30 09:01:52 +020037from .operation import Op
Diego Russoe8a10452020-04-21 17:39:10 +010038from .operation import Operation
Michael McGeagh16895482020-12-14 15:51:20 +000039from .operation import Padding
Fredrik Svedbergd9c2c422020-12-01 16:33:45 +010040from .operation_util import create_avgpool_nop
patrik.gustavssoneeb85152020-12-21 17:10:40 +000041from .shape4d import Shape4D
Fredrik Svedberga0c36242020-06-03 15:43:31 +020042from .softmax import SoftMax
Tim Hall93582962020-09-09 21:58:15 +010043from .tensor import check_quantized_tens_scaling_equal
Michael McGeaghc5b549b2020-08-07 11:54:28 +010044from .tensor import create_const_tensor
45from .tensor import create_reshape_tensor
Charles Xu9a03fdf2020-07-02 15:12:40 +020046from .tensor import QuantizationParameters
Diego Russoe8a10452020-04-21 17:39:10 +010047from .tensor import Tensor
Michael McGeagh7a6f8432020-12-02 15:29:22 +000048from .tflite_mapping import optype_to_builtintype
Tim Hall79d07d22020-04-27 18:20:16 +010049
Michael McGeaghf3e3ad72020-12-02 12:39:03 +000050passthrough_nodes = (Op.Identity,)
Tim Hall79d07d22020-04-27 18:20:16 +010051
Michael McGeaghf3e3ad72020-12-02 12:39:03 +000052memory_only_ops = (Op.Reshape,)
Michael McGeagh11b0bdb2020-09-08 11:07:35 +010053
Tim Hall79d07d22020-04-27 18:20:16 +010054
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +020055def remove_passthrough_tensor(tens, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +010056 if len(tens.ops) == 1 and tens.ops[0].type in passthrough_nodes:
57 assert len(tens.ops[0].inputs) == 1
58 tens = tens.ops[0].inputs[0]
59 return tens
60
61
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +020062def rewrite_concat(tens, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +020063 if len(tens.ops) == 1 and tens.ops[0].type.is_concat_op():
Tim Hall79d07d22020-04-27 18:20:16 +010064 concat_op = tens.ops[0]
65 if tens != concat_op.outputs[0]:
66 return tens # don't attempt to rewrite the min/max outputs of QuantizedConcat
67
68 # Not supported so leave it and run on CPU
69 if not concat_op.run_on_npu:
70 return tens
71
72 inputs, axis = concat_op.get_concat_inputs_axis()
73
74 tens.ops = []
75 offset = 0
76 for idx, inp in enumerate(inputs):
Patrik Gustavsson3d737172020-12-22 10:40:51 +010077 if axis >= 0:
78 axis_4D = axis + (4 - len(inp.shape))
79 else:
80 axis_4D = axis
Louis Verhaardaee5d752020-09-30 09:01:52 +020081 new_op = Operation(Op.ConcatSliceWrite, concat_op.name + str(idx))
Tim Hall79d07d22020-04-27 18:20:16 +010082 new_op.inputs = [inp]
83 new_op.outputs = [tens]
Patrik Gustavsson3d737172020-12-22 10:40:51 +010084 new_op.attrs["concat_axis"] = axis_4D
Tim Hall79d07d22020-04-27 18:20:16 +010085 new_op.attrs["concat_start"] = offset
86 offset += inp.shape[axis]
87 new_op.attrs["concat_end"] = offset
88 new_op.run_on_npu = True
89 tens.ops.append(new_op)
Tim Halle6ccd872020-11-09 16:46:37 +000090 DebugDatabase.add_optimised(concat_op, new_op)
patrik.gustavssoneeb85152020-12-21 17:10:40 +000091 new_op.set_ifm_ofm_shapes()
Tim Hall79d07d22020-04-27 18:20:16 +010092 assert tens.shape[axis] == offset
93
Patrik Gustavsson29d568e2020-08-18 10:11:21 +020094 # If axis corresponds to C-dimension, NHCWB16 can only be used in the output if all the concat_start's are a
95 # multiple of 16. This as, it is only then the address offset for the ofm, for all operations, will be 16 byte
96 # aligned. For other values of axis the address offsets will be 16 byte aligned, as they are all based on c = 0
Patrik Gustavsson458a2082020-08-13 13:41:05 +020097 # and those addresses are always 16 byte aligned due to the NHCWB16 format.
Patrik Gustavsson6c97e9a2020-09-23 11:02:18 +020098 if axis == -1 or axis == (len(tens.shape) - 1):
Patrik Gustavsson458a2082020-08-13 13:41:05 +020099 for op in tens.ops:
100 if op.attrs["concat_start"] % 16 != 0:
101 tens.avoid_NHCWB16 = True
102 break
103
Tim Hall79d07d22020-04-27 18:20:16 +0100104 return tens
105
106
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200107def rewrite_split(tens, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +0100108
Louis Verhaardaee5d752020-09-30 09:01:52 +0200109 if len(tens.ops) == 1 and tens.ops[0].type.is_split_op():
Tim Hall79d07d22020-04-27 18:20:16 +0100110 split_op = tens.ops[0]
111
112 # Not supported so leave it and run on CPU
113 if not split_op.run_on_npu:
114 return tens
115
116 inp, outputs, axis, offset_start, offset_end = split_op.get_split_inputs_axis()
117
118 tens.ops = []
Louis Verhaardaee5d752020-09-30 09:01:52 +0200119 new_op = Operation(Op.SplitSliceRead, split_op.name)
Tim Hall79d07d22020-04-27 18:20:16 +0100120 new_op.inputs = [inp]
Tim Hall79d07d22020-04-27 18:20:16 +0100121
122 # For Split the offset cannot be extracted from the tensor so it has to
123 # be calculated from the index of the output tensor
Diego Russoea6111a2020-04-14 18:41:58 +0100124 if axis is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100125 # Get the start and end of the split
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100126 offset_start = [0] * 4
127 for idx, out in enumerate(outputs):
Tim Hall79d07d22020-04-27 18:20:16 +0100128 if out == tens:
129 break
Patrik Gustavsson3d737172020-12-22 10:40:51 +0100130 if axis >= 0:
131 axis_4D = axis + (4 - len(out.shape))
132 else:
133 axis_4D = axis
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000134
135 offset_start[axis_4D] += split_op.ofm_shapes[idx].get_dim(axis_4D)
Tim Hall79d07d22020-04-27 18:20:16 +0100136
Patrik Gustavssoneebb1c22020-08-18 15:03:04 +0200137 # If start offset is not a multiple of 16 in the C-dimension, NHCWB16 need to be avoided in the input
138 if (offset_start[-1] % 16) != 0:
139 inp.avoid_NHCWB16 = True
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100140 else:
141 offset_start = full_shape(4, offset_start, 0)
Tim Hall79d07d22020-04-27 18:20:16 +0100142
143 new_op.attrs["split_start"] = offset_start
Tim Hall79d07d22020-04-27 18:20:16 +0100144 new_op.run_on_npu = True
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100145 new_op.set_output_tensor(tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000146 new_op.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000147 DebugDatabase.add_optimised(split_op, new_op)
Tim Hall79d07d22020-04-27 18:20:16 +0100148
149 return tens
150
151
152def needed_total_padding(input_size, stride, filter_size):
153 out_size = (input_size + stride - 1) // stride
154 needed_input = (out_size - 1) * stride + filter_size
155 total_padding = max(0, needed_input - input_size)
156 return total_padding
157
158
Louis Verhaardae2d5532020-12-11 17:19:54 +0100159def calc_padding_and_skirt(padding_type, kernel_size, stride, input_dims, explicit_padding):
Tim Hall79d07d22020-04-27 18:20:16 +0100160 ypad = needed_total_padding(int(input_dims[1]), int(stride[1]), int(kernel_size[0]))
161 xpad = needed_total_padding(int(input_dims[2]), int(stride[2]), int(kernel_size[1]))
Michael McGeagh16895482020-12-14 15:51:20 +0000162 if padding_type == Padding.SAME:
Tim Hall79d07d22020-04-27 18:20:16 +0100163 left_pad = (xpad + 0) // 2
164 right_pad = (xpad + 1) // 2
165 top_pad = (ypad + 0) // 2
166 bottom_pad = (ypad + 1) // 2
Michael McGeagh16895482020-12-14 15:51:20 +0000167 elif padding_type == Padding.VALID:
Tim Hall79d07d22020-04-27 18:20:16 +0100168 left_pad = 0
169 right_pad = 0
170 top_pad = 0
171 bottom_pad = 0
Louis Verhaardae2d5532020-12-11 17:19:54 +0100172 elif padding_type == Padding.EXPLICIT:
173 # Padding is specified in a PAD operator which has been bypassed.
174 # The top and left padding are taken from the PAD; bottom and right are calculated.
175 top_pad, left_pad, _, _ = explicit_padding
176 bottom_pad = ypad - top_pad
177 right_pad = xpad - left_pad
Tim Hall79d07d22020-04-27 18:20:16 +0100178 else:
Michael McGeagh16895482020-12-14 15:51:20 +0000179 raise UnsupportedFeatureError(f"Unknown padding")
Tim Hall79d07d22020-04-27 18:20:16 +0100180 padding = (top_pad, left_pad, bottom_pad, right_pad)
181 skirt = (top_pad, left_pad, ypad - top_pad, xpad - left_pad)
182 return padding, skirt
183
Tim Hallc30f4952020-06-15 20:47:35 +0100184
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200185def calc_upscaled_padding_and_skirt(padding_type, kernel_size, stride, input_dims, upscaling_factor):
186 kernel_height, kernel_width = kernel_size[0], kernel_size[1]
Michael McGeagh16895482020-12-14 15:51:20 +0000187 if padding_type == Padding.SAME:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200188 ypad = needed_total_padding(int(input_dims[1]) * upscaling_factor, int(stride[1]), int(kernel_height))
189 xpad = needed_total_padding(int(input_dims[2]) * upscaling_factor, int(stride[2]), int(kernel_width))
Jacob Bohlind47cc272020-08-24 11:42:14 +0200190 right_pad = max(((xpad + 1) // upscaling_factor) - 1, 0)
191 bottom_pad = max(((ypad + 1) // upscaling_factor) - 1, 0)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200192 left_pad = max(kernel_width - 1 - right_pad, 0)
193 top_pad = max(kernel_height - 1 - bottom_pad, 0)
Michael McGeagh16895482020-12-14 15:51:20 +0000194 elif padding_type == Padding.VALID:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200195 right_pad = max(kernel_width - 2, 0)
196 bottom_pad = max(kernel_height - 2, 0)
197 left_pad = kernel_width - 1
198 top_pad = kernel_height - 1
Jacob Bohlincf7da102020-05-20 09:03:40 +0200199 else:
Michael McGeagh16895482020-12-14 15:51:20 +0000200 raise UnsupportedFeatureError(f"Unknown padding")
Jacob Bohlincf7da102020-05-20 09:03:40 +0200201 padding = (top_pad, left_pad, bottom_pad, right_pad)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200202 skirt = padding
Jacob Bohlincf7da102020-05-20 09:03:40 +0200203 return padding, skirt
204
Tim Hall79d07d22020-04-27 18:20:16 +0100205
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200206def fixup_conv2d_backprop(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200207 if op.type == Op.Conv2DBackpropInput:
Tim Hall79d07d22020-04-27 18:20:16 +0100208 # flip the inputs
209 op.inputs[0], op.inputs[2] = op.inputs[2], op.inputs[0]
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000210 op.set_ifm_ofm_shapes()
Louis Verhaardaee5d752020-09-30 09:01:52 +0200211 op.type = Op.Conv2DBackpropInputSwitchedBias
Jacob Bohlincf7da102020-05-20 09:03:40 +0200212
213 # Update strides
Tim Hallc30f4952020-06-15 20:47:35 +0100214 op.attrs.update({"stride_w": 1, "stride_h": 1, "strides": (1, 1, 1, 1)})
Tim Hall79d07d22020-04-27 18:20:16 +0100215
216 return op
217
218
Charles Xu9a03fdf2020-07-02 15:12:40 +0200219# Convert the op to an elementwise add
220def convert_resizebilinear_1x1_to_add(op):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200221 op.type = Op.Add
Charles Xu9a03fdf2020-07-02 15:12:40 +0200222 op.name = op.name + "_add"
Charles Xu9a03fdf2020-07-02 15:12:40 +0200223 op.attrs["resizebilinear"] = True
224 # Create an input tensor filled with zeros
225 shape = op.outputs[0].shape
226 tens = Tensor(shape, op.inputs[0].dtype, op.inputs[1].name + "_add")
227 tens.values = np.zeros(shape)
228 tens.quant_values = np.zeros(shape, np.uint8)
229 tens.quantization = QuantizationParameters(0.0, 255.0)
230 tens.quantization.scale_f32 = 1.0
231 tens.quantization.zero_point = 0
232 tens.consumer_list = [op]
233 tens_op = op.inputs[1].ops[0]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100234 tens_op.set_output_tensor(tens)
Charles Xu9a03fdf2020-07-02 15:12:40 +0200235 # Set the add inputs
236 op.inputs[1] = op.inputs[0]
237 op.inputs[0] = tens
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000238 op.set_ifm_ofm_shapes()
Charles Xu9a03fdf2020-07-02 15:12:40 +0200239
240 return op
241
242
Charles Xu87c13502020-08-06 12:17:26 +0200243# Convert ResizeBilinear to a number of 2x2 pool ops
244def convert_resizebilinear_to_2x2_pool(op):
245 count = 0
246 pre_op = op
247 outputs = op.outputs
248
249 op.attrs.update({"strides": (1, 1, 1, 1), "ksize": (1, 2, 2, 1)})
250 if op.attrs["align_corners"]:
251 shape_modifier = 1
Michael McGeagh16895482020-12-14 15:51:20 +0000252 op.attrs["padding"] = Padding.VALID
Charles Xu87c13502020-08-06 12:17:26 +0200253 else:
254 shape_modifier = 0
Michael McGeagh16895482020-12-14 15:51:20 +0000255 op.attrs["padding"] = Padding.SAME
Charles Xu87c13502020-08-06 12:17:26 +0200256 op.inputs[0].resampling_mode = resampling_mode.NEAREST
257
258 upscaled_shape = np.array(op.inputs[0].shape[1:3])
259 out_shape = np.array(op.outputs[0].shape[1:3])
260 if (upscaled_shape == upscaled_shape * 2 - shape_modifier).all():
261 return op
262
263 while (upscaled_shape < out_shape).all():
264 if count == 0:
265 scaled_op = pre_op
266 else:
267 scaled_op = op.clone("_{}".format(count))
268 scaled_op.inputs[0] = pre_op.outputs[0]
269
270 upscaled_shape = upscaled_shape * 2 - shape_modifier
271
272 if (upscaled_shape == out_shape).all():
273 scaled_op.outputs = outputs
274 scaled_op.outputs[0].ops = [scaled_op]
275 else:
276 shape = outputs[0].shape.copy()
277 shape[1:3] = upscaled_shape[0:2]
278 out_tens = Tensor(shape, DataType.int16, "{}_{}".format(op.outputs[0].name, count))
279 out_tens.quantization = op.outputs[0].quantization.clone()
280 out_tens.quantization.quant_min = np.iinfo(np.int16).min
281 out_tens.quantization.quant_max = np.iinfo(np.int16).max
282 scaled_op.set_output_tensor(out_tens)
283 pre_op = scaled_op
284 count += 1
285
286 # Setup the scale value
287 if scaled_op.inputs[0].dtype.bits == 8 and scaled_op.outputs[0].dtype.bits == 16:
288 scaled_op.attrs["rescale"] = 128
289 elif scaled_op.inputs[0].dtype.bits == 16 and scaled_op.outputs[0].dtype.bits == 8:
290 scaled_op.attrs["rescale"] = 1 / 128
291 elif "rescale" in scaled_op.attrs:
292 del scaled_op.attrs["rescale"]
Patrik Gustavssoncc6915c2020-12-22 09:16:50 +0100293 scaled_op.set_ifm_ofm_shapes()
Charles Xu87c13502020-08-06 12:17:26 +0200294
295 return op
296
297
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200298def fixup_resizebilinear(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200299 if op.type == Op.ResizeBilinear and op.run_on_npu:
Charles Xu87c13502020-08-06 12:17:26 +0200300 if op.inputs[0].shape == op.outputs[0].shape:
Charles Xu36ffaf32020-08-05 15:40:44 +0200301 # Bypass nop resizebilinear
302 op.inputs = op.inputs[:1]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200303 op.type = Op.Identity
Charles Xu87c13502020-08-06 12:17:26 +0200304 elif op.inputs[0].shape[1] == 1 and op.inputs[0].shape[2] == 1:
305 convert_resizebilinear_1x1_to_add(op)
306 else:
307 convert_resizebilinear_to_2x2_pool(op)
Charles Xu9a03fdf2020-07-02 15:12:40 +0200308
309 return op
310
311
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200312def convert_nop_split_to_identity(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200313 if op.type == Op.Split and op.attrs.get("num_splits") == 1:
Dwight Lidmanc3862c22020-09-14 15:22:33 +0200314 # the list comprehension should return a list with a single tensor
315 # if it shouldn't, remove_passthrough_tensor will fail appropriately
316 op.inputs = [i for i in op.inputs if i.shape == op.outputs[0].shape]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200317 op.type = Op.Identity
Dwight Lidmanc3862c22020-09-14 15:22:33 +0200318 return op
319
320
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200321def fixup_fully_connected_input(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200322 if op.type == Op.FullyConnected:
Tim Hall79d07d22020-04-27 18:20:16 +0100323 inp = op.inputs[0]
324 weights = op.inputs[1]
325
326 n_in_elems = weights.shape[-2]
327 elms = inp.elements()
328 batch_size = elms // n_in_elems
329 assert batch_size * n_in_elems == elms
330
331 desired_shape = [batch_size, n_in_elems]
332 if inp.shape != desired_shape:
333 # mismatch, insert a reshape to fix this.
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200334 op.set_input_tensor(create_reshape_tensor(inp, desired_shape), 0)
Tim Hall79d07d22020-04-27 18:20:16 +0100335
336 return op
337
338
Diqing Zhong94457b12020-12-09 15:22:40 +0100339def convert_batched_fc_shape(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200340 if op.type == Op.FullyConnected:
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200341 ifm = op.inputs[0]
342 ofm = op.outputs[0]
343 # Check if the FC is 2D and first dimension indicates batching
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100344 # TOD0 op.ifm_shape[0] > 1 is enough when refactory is complete
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000345 if len(ifm.shape) == len(ofm.shape) == 2 and ifm.shape[0] > 1 and op.ifm_shapes[0].batch > 1:
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200346 n = ifm.shape[0]
347 batching_split = {4: (2, 2), 8: (2, 4), 16: (4, 4)}
348 h, w = batching_split.get(n, (1, n))
349
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200350 prev_op = ifm.ops[0]
351 desired_shape = [1, h, w, ifm.shape[-1]]
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000352 op.ifm_shapes[0] = Shape4D(desired_shape)
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100353
Louis Verhaardaee5d752020-09-30 09:01:52 +0200354 if len(ifm.consumer_list) == 1 and prev_op is not None and prev_op.type == Op.Reshape:
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200355 # There is a preceding Reshape
356 # Compare input of prev_op and input of op, to see if prev_op can be removed
357 ifm_prev_op = prev_op.inputs[0]
Tim Hall89567612020-10-27 11:57:57 +0000358 if ifm_prev_op.shape == ifm.shape and check_quantized_tens_scaling_equal(ifm_prev_op, ifm):
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200359 # prev_op can be removed
360 op.set_input_tensor(ifm_prev_op, 0)
361 else:
362 op.inputs[0].set_all_shapes(desired_shape)
363 prev_op.set_input_tensor(
364 create_const_tensor(prev_op.inputs[1].name, [1], DataType.int32, desired_shape), 1
365 )
366 prev_op.attrs["new_shape"] = desired_shape
367 else:
368 # Add reshape op to the input if there is no preceding reshape
369 ifm.consumer_list.remove(op)
370 op.set_input_tensor(create_reshape_tensor(ifm, desired_shape), 0)
371
372 # Reshape Weights to be 4D. IO becomes HWIO
373 weight_tensor = op.inputs[1]
374 weight_tensor.quant_values = np.expand_dims(np.expand_dims(weight_tensor.quant_values, axis=0), axis=0)
375 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
376
377 desired_shape = [1, h, w, ofm.shape[-1]]
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000378 op.ofm_shapes[0] = Shape4D(desired_shape)
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100379
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200380 if (
381 len(ofm.consumer_list) == 1
382 and ofm.consumer_list[0] is not None
Louis Verhaardaee5d752020-09-30 09:01:52 +0200383 and ofm.consumer_list[0].type == Op.Reshape
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200384 ):
385 # There is a subsequent Reshape
386 # Compare desired shape and output of consumer op, to see if consumer op can be removed
387 ofm_cons_op = ofm.consumer_list[0].outputs[0]
Tim Hall93582962020-09-09 21:58:15 +0100388 if desired_shape == ofm_cons_op.shape and check_quantized_tens_scaling_equal(ofm, ofm_cons_op):
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200389 op.outputs[0] = ofm_cons_op
390 op.outputs[0].ops = [op]
391 else:
392 op.outputs[0].set_all_shapes(desired_shape)
393 else:
Diqing Zhong94457b12020-12-09 15:22:40 +0100394 # Add reshape op to the output
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200395 op.set_output_tensor(create_reshape_tensor(ofm, desired_shape, False))
396 return op
397
398
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200399def fixup_pack_input(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200400 if op.type == Op.Pack:
Tim Hall79d07d22020-04-27 18:20:16 +0100401 # Pack is also referred to as Stack
402 # Requires the rewrite_concat function to be called on the op afterwards
403 axis = int(op.attrs["axis"])
404 desired_shape = op.inputs[0].shape[:axis] + [1] + op.inputs[0].shape[axis:]
405
406 # Construct 1 shape tensor to be used by all inserted reshape ops
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100407 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, desired_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100408
409 for idx, inp in enumerate(op.inputs):
Tim Hall79d07d22020-04-27 18:20:16 +0100410 reshape_out = inp.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100411 reshape_out.set_all_shapes(desired_shape)
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100412
Louis Verhaardaee5d752020-09-30 09:01:52 +0200413 reshape_op = Operation(Op.Reshape, "{}{}_reshape".format(op.name, idx))
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100414 reshape_op.attrs["new_shape"] = desired_shape
415 reshape_op.inputs = [inp, new_shape_tens]
416 reshape_op.set_output_tensor(reshape_out)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000417 reshape_op.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000418 DebugDatabase.add_optimised(op, reshape_op)
Tim Hall79d07d22020-04-27 18:20:16 +0100419
420 op.inputs[idx] = reshape_out
421
Louis Verhaardaee5d752020-09-30 09:01:52 +0200422 op.type = Op.PackReshaped
Tim Hall79d07d22020-04-27 18:20:16 +0100423
424 return op
425
426
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200427def unfuse_activation_function(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200428 if op.type == Op.ConcatTFLite and op.run_on_npu and op.activation is not None:
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100429 act_op = Operation(op.activation.op_type, op.name + op.activation.op_type.name)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200430 op.activation = None
Fredrik Svedberg0f98b362020-09-29 10:00:39 +0200431 out_tens = op.outputs[0]
432 intermediate_tens = out_tens.clone("_act_intermediate")
433 act_op.set_output_tensor(out_tens)
434 act_op.add_input_tensor(intermediate_tens)
435 op.set_output_tensor(intermediate_tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000436 act_op.set_ifm_ofm_shapes()
Fredrik Svedberg0f98b362020-09-29 10:00:39 +0200437
438 return op
439
Louis Verhaard8912c532020-09-30 12:11:49 +0200440
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100441def fixup_stridedslice_output(tens, arch, nng):
442 op = tens.ops[0]
Dwight Lidman73320a42020-11-05 10:34:41 +0100443 if op.run_on_npu and op.type == Op.StridedSlice:
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100444 reshape_input_shape = tens.shape
445 new_axis_mask = op.attrs["new_axis_mask"]
446 shrink_axis_mask = op.attrs["shrink_axis_mask"]
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100447
Dwight Lidman73320a42020-11-05 10:34:41 +0100448 if shrink_axis_mask != 0:
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100449 n = 0
450 axis = 0
451 while shrink_axis_mask:
452 prev_mask = shrink_axis_mask
453 n += 1
454 shrink_axis_mask &= shrink_axis_mask - 1
455 axis = int(math.log2(prev_mask - shrink_axis_mask))
456 reshape_input_shape = reshape_input_shape[:axis] + [1] + reshape_input_shape[axis:]
457
458 assert len(tens.shape) == (len(op.inputs[0].shape) - n)
459 op.attrs["shrink_axis_mask"] = 0
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100460 elif new_axis_mask != 0:
461 n = 0
462 axis = 0
463 while new_axis_mask:
464 prev_mask = new_axis_mask
465 n += 1
466 new_axis_mask &= new_axis_mask - 1
467 axis = int(math.log2(prev_mask - new_axis_mask))
468 reshape_input_shape = reshape_input_shape[:axis] + reshape_input_shape[(axis + 1) :]
469 new_axis_mask >>= 1
470
471 assert len(tens.shape) == (len(op.inputs[0].shape) + n)
472 op.attrs["new_axis_mask"] = 0
Dwight Lidmandda21af2020-11-11 15:44:57 +0100473 else:
474 # Equal Rank StridedSlice, no need to insert reshape
475 return tens
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100476
477 # Construct 1 shape tensor to be used by all inserted reshape ops
478 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, tens.shape)
479
480 for idx, out_tens in enumerate(op.outputs):
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000481 op.ofm_shapes[idx] = Shape4D(new_shape_tens.shape)
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100482 reshape_in = out_tens.clone("_reshaped")
483 reshape_in.set_all_shapes(reshape_input_shape)
484 reshape_in.ops = [op]
485
486 reshape_op = Operation(Op.Reshape, "{}{}_reshape".format(op.name, idx))
487 reshape_op.attrs["new_shape"] = reshape_input_shape
488 reshape_op.inputs = [reshape_in, new_shape_tens]
489 reshape_op.set_output_tensor(out_tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000490 reshape_op.set_ifm_ofm_shapes()
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100491
492 op.outputs[idx] = reshape_in
493
494 return tens
495
496
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200497def fixup_unpack_output(tens, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +0100498 op = tens.ops[0]
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100499 if op.run_on_npu and op.type == Op.Unpack:
Tim Hall79d07d22020-04-27 18:20:16 +0100500 # Unpack is also referred to as Unstack
501 # Requires the rewrite_split function to be called on the op afterwards
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100502 axis = int(op.attrs["axis"])
503 op.type = Op.UnpackReshaped
504 reshape_input_shape = tens.shape[:axis] + [1] + tens.shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100505
506 # Construct 1 shape tensor to be used by all inserted reshape ops
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100507 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, tens.shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100508
509 for idx, out_tens in enumerate(op.outputs):
Tim Hall79d07d22020-04-27 18:20:16 +0100510 reshape_in = out_tens.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100511 reshape_in.set_all_shapes(reshape_input_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100512 reshape_in.ops = [op]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100513
Louis Verhaardaee5d752020-09-30 09:01:52 +0200514 reshape_op = Operation(Op.Reshape, "{}{}_reshape".format(op.name, idx))
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100515 reshape_op.attrs["new_shape"] = reshape_input_shape
Tim Hall79d07d22020-04-27 18:20:16 +0100516 reshape_op.inputs = [reshape_in, new_shape_tens]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100517 reshape_op.set_output_tensor(out_tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000518 reshape_op.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000519 DebugDatabase.add_optimised(op, reshape_op)
Tim Hall79d07d22020-04-27 18:20:16 +0100520
521 op.outputs[idx] = reshape_in
Tim Hall79d07d22020-04-27 18:20:16 +0100522 return tens
523
524
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200525def add_padding_fields(op, arch, nng):
Jacob Bohlin90033f32020-08-28 15:45:44 +0200526 if op.run_on_npu:
527 if "padding" in op.attrs:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200528 if op.type.is_conv2d_op() or op.type.is_depthwise_conv2d_op():
Jacob Bohlin90033f32020-08-28 15:45:44 +0200529 kernel_size = op.inputs[1].shape[:2]
530 input_shape = op.inputs[0].shape
Louis Verhaardaee5d752020-09-30 09:01:52 +0200531 elif op.type.is_pool_op() or op.type.npu_block_type == NpuBlockType.ReduceSum:
Jacob Bohlin90033f32020-08-28 15:45:44 +0200532 kernel_size = op.attrs["ksize"][1:3]
533 input_shape = op.inputs[0].shape
Jacob Bohlin90033f32020-08-28 15:45:44 +0200534 else:
Michael McGeagh7a6f8432020-12-02 15:29:22 +0000535 raise UnsupportedFeatureError(f"Unknown operation that uses padding: {optype_to_builtintype(op.type)}")
Tim Hall79d07d22020-04-27 18:20:16 +0100536
Louis Verhaardaee5d752020-09-30 09:01:52 +0200537 if op.type == Op.Conv2DBackpropInputSwitchedBias:
Jacob Bohlin90033f32020-08-28 15:45:44 +0200538 upscaling_factor = op.outputs[0].shape[1] // input_shape[1]
539 padding, skirt = calc_upscaled_padding_and_skirt(
540 op.attrs["padding"], kernel_size, op.attrs["strides"], input_shape, upscaling_factor
541 )
542 else:
543 dilation_h, dilation_w = op.get_dilation_h_w()
544 dilated_kernel_size = [dilation_h * (kernel_size[0] - 1) + 1, dilation_w * (kernel_size[1] - 1) + 1]
545 padding, skirt = calc_padding_and_skirt(
Louis Verhaardae2d5532020-12-11 17:19:54 +0100546 op.attrs["padding"],
547 dilated_kernel_size,
548 op.attrs["strides"],
549 input_shape,
550 op.attrs.get("explicit_padding"),
Jacob Bohlin90033f32020-08-28 15:45:44 +0200551 )
Jacob Bohlincf7da102020-05-20 09:03:40 +0200552
Jacob Bohlin90033f32020-08-28 15:45:44 +0200553 op.attrs["explicit_padding"] = padding
554 op.attrs["skirt"] = skirt
Jacob Bohlincf7da102020-05-20 09:03:40 +0200555
Tim Hall79d07d22020-04-27 18:20:16 +0100556 return op
557
558
Tim Hall79d07d22020-04-27 18:20:16 +0100559# Check if the op can be reordered
560def get_prepend_op(op):
561 inp = op.inputs[0]
562 # The op should be reordered between prev_op and prep_op
563 prev_op = inp.ops[-1]
564 prep_op = None
565 while prev_op.type in memory_only_ops and len(prev_op.outputs) == 1 and len(prev_op.outputs[0].consumers()) == 1:
566 prep_op = prev_op
567 inp = prev_op.inputs[0]
568 prev_op = inp.ops[-1]
Diego Russoea6111a2020-04-14 18:41:58 +0100569 if prev_op is not None and len(prev_op.outputs) == 1 and len(prev_op.outputs[0].consumers()) == 1:
Tim Hall79d07d22020-04-27 18:20:16 +0100570 return prep_op
571
572 return None
573
574
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200575def convert_depthwise_to_conv(op, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +0100576 # Depthwise is equivalent to a single conv2d if the ifm depth is 1 and
577 # the ofm depth equals the depth multipler.
578 # If those conditions are true, then we can perform a simple
579 # switch of the operator type (and weight order)
580
Louis Verhaardaee5d752020-09-30 09:01:52 +0200581 if op.type == Op.DepthwiseConv2DBias and (op.attrs["depth_multiplier"] != 1):
Tim Hall79d07d22020-04-27 18:20:16 +0100582 ifm_tensor = op.inputs[0]
583 weight_tensor = op.inputs[1]
584 ofm_tensor = op.outputs[0]
585 if (ifm_tensor.shape[3] == 1) and (ofm_tensor.shape[3] == op.attrs["depth_multiplier"]):
586 # Change op type to Conv2d
Louis Verhaardaee5d752020-09-30 09:01:52 +0200587 op.type = Op.Conv2DBias
Tim Hall79d07d22020-04-27 18:20:16 +0100588 del op.attrs["channel_multiplier"]
589 del op.attrs["depth_multiplier"]
590
591 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100592 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Tim Hall79d07d22020-04-27 18:20:16 +0100593 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200594 raise UnsupportedFeatureError(
Michael McGeagh7a6f8432020-12-02 15:29:22 +0000595 f"Unsupported 'DEPTHWISE_CONV_2D' with depth_multiplier = {op.attrs['depth_multiplier']},",
596 f" ifm channels = {ifm_tensor.shape[3]}, ofm channels = {ofm_tensor.shape[3]}",
Tim Hall79d07d22020-04-27 18:20:16 +0100597 )
Tim Halle6ccd872020-11-09 16:46:37 +0000598 DebugDatabase.add_optimised(op, op)
Tim Hall79d07d22020-04-27 18:20:16 +0100599 return op
600
601
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200602def reorder_depthwise_weights(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200603 if op.type.is_depthwise_conv2d_op():
Jacob Bohline843d332020-06-23 12:12:56 +0200604 weight_tensor = op.inputs[1]
605 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100606 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Jacob Bohline843d332020-06-23 12:12:56 +0200607 weight_tensor.weight_transpose_depthwise = True
608
609 return op
610
611
Diqing Zhong016b8272020-12-16 16:46:06 +0100612def optimise_strided_conv(op, arch, nng):
613 stride_x, stride_y = op.get_kernel_stride()
614 ifm_tensor, _, weight_tensor, _ = op.get_ifm_ifm2_weights_ofm()
615
616 if (
617 op.type == Op.Conv2DBias
618 and op.op_index == 0
619 and stride_x == 2
620 and len(ifm_tensor.shape) == 4
621 and ifm_tensor.shape[3] <= 4
622 and ifm_tensor.shape[2] % 2 == 0
623 and weight_tensor is not None
624 and weight_tensor.shape[1] >= 2
625 ):
626 # IFM
627 ifm_reshaped = create_reshape_tensor(
628 ifm_tensor, [ifm_tensor.shape[0], ifm_tensor.shape[1], ifm_tensor.shape[2] // 2, ifm_tensor.shape[3] * 2]
629 )
630 op.set_input_tensor(ifm_reshaped, 0)
631
632 # Weights
633 weight_shape = weight_tensor.shape
634 if weight_shape[1] % 2 != 0:
635 weight_shape[1] = weight_shape[1] + 1
636 padded_array = np.zeros(weight_shape)
637 for i in range(weight_shape[0]):
638 padded_array[i] = np.vstack(
639 [
640 weight_tensor.quant_values[i],
641 np.full((1, weight_shape[2], weight_shape[3]), weight_tensor.quantization.zero_point),
642 ]
643 )
644 weight_tensor.quant_values = padded_array
645 weight_shape[1] //= 2
646 weight_shape[2] *= 2
647 weight_tensor.quant_values = np.reshape(weight_tensor.quant_values, weight_shape)
648 weight_tensor.set_all_shapes(weight_shape)
649 # If multiple copies of the weights are used, we could avoid
650 # them having the same address by changing the value_id
651 weight_tensor.value_id = uuid.uuid4()
652
653 # Strides
654 stride_x = 1
655 op.attrs.update({"stride_w": stride_x, "stride_h": stride_y, "strides": (1, stride_y, stride_x, 1)})
656
657 op.set_ifm_ofm_shapes()
658
659 return op
660
661
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200662def convert_conv_to_fc(op, arch, nng):
Michael McGeagh8d939c02020-07-29 13:11:43 +0100663 # Conv 1x1 can be equivalent to Fully Connected.
664 # By representing certain convs as fully connected layers, Vela can better determine wether or not to use
665 # caching/double buffering for the weights.
666 # (Weights dont need to be reloaded for convs when IFM H and W are 1)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200667 if op.type == Op.Conv2DBias:
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000668 h = op.ifm_shapes[0].height
669 w = op.ifm_shapes[0].width
Michael McGeagh8d939c02020-07-29 13:11:43 +0100670 kh, kw, _, _ = op.inputs[1].shape
671 if h == 1 and w == 1 and kh == 1 and kw == 1:
672 # Overwrite this op as a Fully Connected Op
673 op.name += "_fc"
Louis Verhaardaee5d752020-09-30 09:01:52 +0200674 op.type = Op.FullyConnected
Michael McGeagh8d939c02020-07-29 13:11:43 +0100675 op.attrs = {
Michael McGeagh8d939c02020-07-29 13:11:43 +0100676 "weights_format": 0,
Michael McGeagh8d939c02020-07-29 13:11:43 +0100677 }
678 # Reshape Weights to be 2D. HWIO becomes just IO (as H and W are 1, they can just be dropped)
679 weight_tensor = op.inputs[1]
680 weight_tensor.quant_values = weight_tensor.quant_values.squeeze(axis=(0, 1))
681 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100682
Michael McGeagh8d939c02020-07-29 13:11:43 +0100683 # The output from a fully connected is expected to be 2D so we need to add a reshape layer to convert it
684 # back to 4D afterwards as the next layer is expecting that shape
685 orig_ofm_tensor = op.outputs[0]
686 # Reshape this ops output to be 2D: {(N*H*W), C} (We know N H and W are all 1 so this becomes {1, C})
687 fc_ofm_tensor = orig_ofm_tensor.clone("_fc")
688 fc_ofm_tensor.set_all_shapes([1, fc_ofm_tensor.shape[-1]])
689 fc_ofm_tensor.ops = [op]
690 # Add a reshape after the new OFM to convert it back to the original 4D shape
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100691 reshape_name = op.name + "_reshape"
692 new_shape_tens = create_const_tensor(reshape_name + "_shape", [1], DataType.int32, orig_ofm_tensor.shape)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200693 reshape_op = Operation(Op.Reshape, reshape_name)
Michael McGeagh8d939c02020-07-29 13:11:43 +0100694 reshape_op.attrs["new_shape"] = orig_ofm_tensor.shape
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100695 reshape_op.inputs = [fc_ofm_tensor, new_shape_tens]
696 reshape_op.set_output_tensor(orig_ofm_tensor)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000697 reshape_op.set_ifm_ofm_shapes()
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100698
Michael McGeagh8d939c02020-07-29 13:11:43 +0100699 # Replace this ops OFM to point to the 2D tensor
700 op.outputs[0] = fc_ofm_tensor
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000701 op.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000702 # Record optimisation in debug database
703 DebugDatabase.add_optimised(op, reshape_op)
704 DebugDatabase.add_optimised(op, op)
Michael McGeagh8d939c02020-07-29 13:11:43 +0100705 return op
706
707
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200708def fixup_relus_with_differing_ifm_ofm_scaling(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200709 if op.run_on_npu and op.type.is_relu_op():
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100710 ifm = op.inputs[0]
711 ofm = op.outputs[0]
712 # Relu with differing IFM and OFM scaling cannot be fused with another primary op
713 # and requires its own to be inserted
Tim Hall93582962020-09-09 21:58:15 +0100714 if not check_quantized_tens_scaling_equal(ifm, ofm):
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100715 # Override this op with its own primary op (avgpool)
716 relu_fused_op = create_avgpool_nop(op.name + "_avgpool")
717 # And fuse the original activation function to it
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100718 relu_fused_op.activation = create_activation_function(op.type)
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100719 # Tidy up and assign the ifm and ofm to the new op
720 ifm.consumer_list.remove(op)
Andreas Nevalainenf3d737e2020-09-25 14:12:43 +0200721
722 # if not 4d, reshape ifm/ofm
723 if len(ifm.shape) < 4:
724 ifm_shaped = create_reshape_tensor(ifm, full_shape(4, ifm.shape, 1))
725 ifm = ifm_shaped
726 if len(ofm.shape) < 4:
727 ofm_shaped = create_reshape_tensor(ofm, full_shape(4, ofm.shape, 1), False)
728 ofm = ofm_shaped
729
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100730 relu_fused_op.add_input_tensor(ifm)
731 relu_fused_op.set_output_tensor(ofm)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000732 relu_fused_op.set_ifm_ofm_shapes()
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100733 op = relu_fused_op
734 return op
735
736
Tim Hall79d07d22020-04-27 18:20:16 +0100737# Reorder activation op if it's after the memory only operations
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200738def fixup_act_reorder(op, arch, nng):
Michael McGeaghf3e3ad72020-12-02 12:39:03 +0000739 if op.type.is_relu_op() or op.type in (Op.Sigmoid, Op.Tanh):
Tim Hall79d07d22020-04-27 18:20:16 +0100740 prep_op = get_prepend_op(op)
Diego Russoea6111a2020-04-14 18:41:58 +0100741 if prep_op is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100742 act_op = op.clone("_reordered")
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100743 act_op.ifm_shapes = list(op.ifm_shapes)
744 act_op.ofm_shapes = list(op.ofm_shapes)
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200745
746 # There is only one input tensor, overwrite it
747 act_op.set_input_tensor(prep_op.inputs[0], 0)
748
Tim Hall79d07d22020-04-27 18:20:16 +0100749 act_op_out = act_op.inputs[0].clone("_acted")
750 act_op_out.quantization = op.outputs[0].quantization.clone()
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100751 act_op.set_output_tensor(act_op_out)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000752 act_op.ifm_shapes[0] = Shape4D(prep_op.inputs[0].shape)
753 act_op.ofm_shapes[0] = Shape4D(act_op_out.shape)
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200754
755 # Update the consumer list
756 act_op_out.consumer_list = op.outputs[0].consumer_list.copy()
757 act_op_out.consumer_list.append(prep_op)
758
Tim Hall79d07d22020-04-27 18:20:16 +0100759 prep_op.inputs[0] = act_op_out
760 prep_op.outputs[0].quantization = act_op_out.quantization.clone()
761
762 # Mark the op so that it will be removed as passthrough later on
Louis Verhaardaee5d752020-09-30 09:01:52 +0200763 op.type = Op.Identity
Tim Halle6ccd872020-11-09 16:46:37 +0000764
765 # Record optimisation in debug database
766 DebugDatabase.add_optimised(op, act_op)
767 DebugDatabase.add_optimised(op, op)
Tim Hall79d07d22020-04-27 18:20:16 +0100768 return op
769
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200770
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200771def fixup_elementwise_with_scalars(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200772 if op.type.is_binary_elementwise_op():
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200773 ifm_tensor, ifm2_tensor, _, _ = op.get_ifm_ifm2_weights_ofm()
Charles Xu78792222020-05-13 10:15:26 +0200774 if ifm2_tensor.shape != [] and ifm_tensor.shape != []:
775 diff = len(ifm_tensor.shape) - len(ifm2_tensor.shape)
776 if diff > 0:
777 ifm2_tensor.shape = full_shape(len(ifm_tensor.shape), ifm2_tensor.shape, 1)
778 elif diff < 0:
779 ifm_tensor.shape = full_shape(len(ifm2_tensor.shape), ifm_tensor.shape, 1)
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200780 elif ifm_tensor.shape == [] and ifm_tensor.quant_values is None:
781 # IFM is marked as a scalar, but is a result of an operation; change it to a shape of size 1
782 ifm_tensor.shape = len(ifm2_tensor.shape) * [1]
783 ifm_tensor.storage_shape = ifm_tensor.shape
784 elif ifm2_tensor.shape == [] and ifm2_tensor.quant_values is None:
785 # IFM2 is marked as a scalar, but is a result of an operation; change it to a shape of size 1
786 ifm2_tensor.shape = len(ifm_tensor.shape) * [1]
787 ifm2_tensor.storage_shape = ifm2_tensor.shape
Charles Xu78792222020-05-13 10:15:26 +0200788 return op
Tim Hall79d07d22020-04-27 18:20:16 +0100789
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200790
Tim Hall4e127762020-05-15 16:05:49 +0100791# Set input/output tensor equivalence to the same id for memory operations
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200792def set_tensor_equivalence(op, arch, nng):
Michael McGeagh11b0bdb2020-09-08 11:07:35 +0100793 if op.type in memory_only_ops:
Tim Hall4e127762020-05-15 16:05:49 +0100794 eid = op.outputs[0].equivalence_id
795 for inp in op.inputs:
796 inp.equivalence_id = eid
797 return op
798
799
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100800def set_ifm_ofm_op_shapes(op, arch, nng):
801 if op.run_on_npu and op.type.needs_shapes():
802 if op.ifm_shapes or op.ofm_shapes:
803 # Shapes already set
804 return op
805 op.set_ifm_ofm_shapes()
806 return op
807
808
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200809def convert_softmax(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200810 if op.type == Op.Softmax and op.run_on_npu:
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200811 softmax = SoftMax(op)
812 op = softmax.get_graph()
813 return op
814
815
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200816def convert_mul_max_to_abs_or_lrelu(op, arch, nng):
Diego Russoea6111a2020-04-14 18:41:58 +0100817 r"""Whenever there is a subgraph with this topology:
Tim Hall79d07d22020-04-27 18:20:16 +0100818
819 Input X For X = -1 or X > 0
820 | \ / This subgraph can be replaced with either
821 | Mul an Abs (if X = -1) or a LeakyReLU (if X > 0)
822 | /
823 Max
824 """
825
Louis Verhaardaee5d752020-09-30 09:01:52 +0200826 if op.type == Op.Maximum:
Tim Hall79d07d22020-04-27 18:20:16 +0100827 # finds the Mul input(s) to the Max
Louis Verhaardaee5d752020-09-30 09:01:52 +0200828 muls = [i for i in op.inputs if i.ops[0].type == Op.Mul]
Tim Hall79d07d22020-04-27 18:20:16 +0100829 if len(muls) == 1:
830 mul = muls[0].ops[0]
831 elif len(muls) == 2:
832 # In the case both inputs are Muls, find the one with the same input as the Max
833 mul = [m for m in muls if len(set(op.inputs + m.ops[0].inputs)) == 1][0].ops[0]
834 else:
835 # No Mul inputs
836 return op
837
838 # make sure the Mul doesn't have any other consumers
Louis Verhaardd7911c42020-08-25 13:36:41 +0200839 mul_ofm = mul.outputs[0]
840 if len(mul_ofm.consumers()) != 1:
Tim Hall79d07d22020-04-27 18:20:16 +0100841 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +0200842 # make sure the Mul doesn't have a fused activation function
843 if mul.activation:
Tim Hall79d07d22020-04-27 18:20:16 +0100844 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +0200845 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +0100846 if ifm is None or ofm is None:
847 return op
848
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200849 if ifm.dtype not in (DataType.uint8, DataType.int8) or ifm.dtype != ofm.dtype:
850 return op
Tim Hall93582962020-09-09 21:58:15 +0100851 if not check_quantized_tens_scaling_equal(ifm, ofm) or not check_quantized_tens_scaling_equal(ifm, mul_ofm):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200852 # rewrite to LeakyRelu currently only makes sense if the quantization is identical
853 return op
Tim Hall79d07d22020-04-27 18:20:16 +0100854
855 # finds the branched input that goes to both the Max and the Mul
856 shared = set(op.inputs) & set(mul.inputs)
857 if len(shared) == 1:
858 shared_in = shared.pop()
859 # find the constant scalar input to the Mul
860 const_tens = (set(mul.inputs) - {shared_in}).pop()
861 # check that it is a scalar
862 if const_tens.shape != []:
863 return op
864 const = const_tens.ops[0]
865 # check that it is a constant
Louis Verhaardaee5d752020-09-30 09:01:52 +0200866 if const.type != Op.Const:
Tim Hall79d07d22020-04-27 18:20:16 +0100867 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200868 # Remove the Mul from the shared input's consumers
869 shared_in.consumer_list.remove(mul)
Tim Hall79d07d22020-04-27 18:20:16 +0100870 else:
871 return op
872
873 val = const.outputs[0].values
874 if val >= 0:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200875 new_op = Op.LeakyRelu
Tim Hall79d07d22020-04-27 18:20:16 +0100876 op.attrs["alpha"] = val
Louis Verhaardd7911c42020-08-25 13:36:41 +0200877 # to produce bit exact results, the alpha is not enough;
878 # save additional scaling info in attr "alpha_scale", to be used as input
879 # to the LUT construction
880 alpha_scalar = const_tens.quant_values - const_tens.quantization.zero_point
881 mul_ifm_scale = np.double(ifm.quantization.scale_f32)
882 mul_ifm2_scale = np.double(const_tens.quantization.scale_f32)
883 mul_ofm_scale = np.double(mul_ofm.quantization.scale_f32)
884 alpha_scale, alpha_shift = scaling.elementwise_mul_scale(mul_ifm_scale, mul_ifm2_scale, mul_ofm_scale)
885 op.attrs["alpha_scaling"] = (alpha_scalar, alpha_scale, alpha_shift)
Tim Hall79d07d22020-04-27 18:20:16 +0100886 elif val == -1:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200887 new_op = Op.Abs
Tim Hall79d07d22020-04-27 18:20:16 +0100888 else:
889 return op
890
Louis Verhaardaee5d752020-09-30 09:01:52 +0200891 op.type = new_op
892 op.name = op.name.replace("Maximum", new_op.name)
893 op.outputs[0].name = op.outputs[0].name.replace("Maximum", new_op.name)
Tim Hall79d07d22020-04-27 18:20:16 +0100894 op.inputs = [shared_in]
Tim Halle6ccd872020-11-09 16:46:37 +0000895
896 # Record optimisation in debug database
897 DebugDatabase.add_optimised(op, op)
898
Tim Hall79d07d22020-04-27 18:20:16 +0100899 return op
900
901
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200902def convert_lrelu_to_mul_max(op, arch):
903 # Converts LeakyRelu to Max(alpha * IFM, identity * IFM)
904 # (the opposite of convert_mul_max_to_abs_or_lrelu)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200905 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +0100906 if ifm is None or ofm is None:
907 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200908
909 # Add multiplication with alpha
Louis Verhaardaee5d752020-09-30 09:01:52 +0200910 mul_alpha = Operation(Op.Mul, op.name + "_mul_alpha")
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200911 mul_alpha.add_input_tensor(ifm)
912 # Create const tensor containing alpha as scalar
913 alpha = op.attrs["alpha"]
914 quantization = ifm.quantization.clone()
915 quantization.min = 0
916 quantization.max = alpha * (quantization.quant_max - quantization.quant_min)
917 quantization.scale_f32 = alpha
918 quantization.zero_point = 0
919 alpha_tens = create_const_tensor(op.name + "_alpha_scalar", [], ifm.dtype, [1], np.int8, quantization=quantization)
920 mul_alpha.add_input_tensor(alpha_tens)
921 fm_alpha = ofm.clone(op.name + "_alpha")
922 mul_alpha.set_output_tensor(fm_alpha)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000923 mul_alpha.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000924 DebugDatabase.add_optimised(op, mul_alpha)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200925
Tim Hall93582962020-09-09 21:58:15 +0100926 if check_quantized_tens_scaling_equal(ifm, ofm):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200927 # No identity multiplication is needed
928 fm_id = ifm
929 else:
930 # Add multiplication with identity
Louis Verhaardaee5d752020-09-30 09:01:52 +0200931 mul_identity = Operation(Op.Mul, op.name + "_mul_identity")
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200932 mul_identity.add_input_tensor(ifm)
933 # Create const tensor containing identity as scalar
934 quantization = ifm.quantization.clone()
935 quantization.min = 0
936 quantization.max = quantization.quant_max - quantization.quant_min
937 quantization.scale_f32 = 1
938 quantization.zero_point = 0
939 identity_tens = create_const_tensor(
940 op.name + "_id_scalar", [], ifm.dtype, [1], np.uint8, quantization=quantization
941 )
942 mul_identity.add_input_tensor(identity_tens)
943 fm_id = ofm.clone(op.name + "_id")
944 mul_identity.set_output_tensor(fm_id)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000945 mul_identity.set_ifm_ofm_shapes()
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100946 DebugDatabase.add_optimised(op, mul_identity)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200947
948 # Convert LeakyRelu to Max, add the results of the multiplication(s) as inputs
Louis Verhaardaee5d752020-09-30 09:01:52 +0200949 op.type = Op.Maximum
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200950 op.name = op.name.replace("LeakyRelu", "Maximum")
951 op.inputs = []
952 ifm.consumer_list.remove(op)
953 op.add_input_tensor(fm_alpha)
954 op.add_input_tensor(fm_id)
Tim Halle6ccd872020-11-09 16:46:37 +0000955
956 DebugDatabase.add_optimised(op, op)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200957 return op
958
959
Louis Verhaard2e186c72020-10-09 10:47:04 +0200960def convert_to_lut(op, lut_values, lut_name):
Louis Verhaardf03bad32020-09-25 08:30:44 +0200961 # Rewrite the operation by Add with scalar 0 + LUT activation
962 ifm = op.inputs[0]
Tim Hall93582962020-09-09 21:58:15 +0100963 if ifm is None:
964 return op
Louis Verhaard58520b92020-08-24 16:45:38 +0200965 assert ifm.dtype.size_in_bytes() == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +0200966 op.type = Op.Add
Louis Verhaard2e186c72020-10-09 10:47:04 +0200967 op.name = op.name + "_lut_" + lut_name
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200968 # Mark as no-op to enable potential fusing optimizations
969 op.attrs["is_nop"] = True
970 # Create an input tensor containing scalar zero
971 quantization = QuantizationParameters(0.0, 255.0)
Louis Verhaardd7911c42020-08-25 13:36:41 +0200972 quantization.scale_f32 = ifm.quantization.scale_f32
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200973 quantization.zero_point = 0
Louis Verhaard2e186c72020-10-09 10:47:04 +0200974 tens = create_const_tensor(op.inputs[0].name + "_scalar0", [], ifm.dtype, [0], np.uint8, quantization=quantization)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200975 op.add_input_tensor(tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000976 op.ifm_shapes.append(Shape4D(tens.shape))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100977
Louis Verhaardf03bad32020-09-25 08:30:44 +0200978 # The LUT must be applied without any preceding rescaling (the LUT itself performs the rescale),
979 # so even if the OFM has a different scale than the IFM, the generated OFM scale instructions
980 # should be the same as the IFM
Louis Verhaardaee5d752020-09-30 09:01:52 +0200981 op.forced_output_quantization = ifm.quantization
Louis Verhaard2e186c72020-10-09 10:47:04 +0200982 lut_tensor = lut.create_lut_tensor(op.name + "_values", lut_values, DataType.int8)
Louis Verhaardf03bad32020-09-25 08:30:44 +0200983 op.set_activation_lut(lut_tensor)
984 return op
985
986
Louis Verhaard2e186c72020-10-09 10:47:04 +0200987def convert_to_lut8(op, fn, fn_name):
Louis Verhaardf03bad32020-09-25 08:30:44 +0200988 # Converts op to a no-op + int8/uint8 LUT which is generated with the given function.
989 # fn is a function(real) -> real
Louis Verhaardaee5d752020-09-30 09:01:52 +0200990 ifm, ofm = op.get_ifm_ofm()
Louis Verhaardf03bad32020-09-25 08:30:44 +0200991 if ifm.dtype not in (DataType.uint8, DataType.int8) or ifm.dtype != ofm.dtype:
992 return op
993 # Generate the LUT
994 ifm_scale = np.double(ifm.quantization.scale_f32)
995 ofm_scale = np.double(ofm.quantization.scale_f32)
996 zp_in = ifm.quantization.zero_point
997 zp_out = ofm.quantization.zero_point
998 values = []
999 ix = range(256) if ifm.dtype == DataType.uint8 else range(-128, 128)
1000 quantized_min = min(ix)
1001 quantized_max = max(ix)
1002 for x in ix:
1003 x_real = ifm_scale * (x - zp_in)
1004 y_real = fn(x_real)
1005 lut_result = round_away_zero(zp_out + y_real / ofm_scale)
1006 lut_result = min(quantized_max, max(quantized_min, lut_result))
1007 values.append(lut_result)
Louis Verhaard2e186c72020-10-09 10:47:04 +02001008 return convert_to_lut(op, values, fn_name)
Louis Verhaardf03bad32020-09-25 08:30:44 +02001009
1010
1011def convert_lrelu_to_lut(op, arch):
Louis Verhaardaee5d752020-09-30 09:01:52 +02001012 ifm, ofm = op.get_ifm_ofm()
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001013 # Generate the LUT
Louis Verhaardd7911c42020-08-25 13:36:41 +02001014 alpha = op.attrs["alpha"]
1015 ifm_scale = np.double(ifm.quantization.scale_f32)
1016 ofm_scale = np.double(ofm.quantization.scale_f32)
1017 zp_in = ifm.quantization.zero_point
1018 zp_out = ofm.quantization.zero_point
1019 identity_scale, identity_shift = scaling.elementwise_mul_scale(ifm_scale, 1, ofm_scale)
1020 alpha_scalar = 1
1021 alpha_scale, alpha_shift = scaling.elementwise_mul_scale(ifm_scale, alpha, ofm_scale)
1022 if "alpha_scaling" in op.attrs:
1023 # The LeakyRelu was the result from convert_mul_max_to_abs_or_lrelu
1024 alpha_scalar, alpha_scale, alpha_shift = op.attrs["alpha_scaling"]
1025 values = []
Louis Verhaard58520b92020-08-24 16:45:38 +02001026 ix = range(256) if ifm.dtype == DataType.uint8 else range(-128, 128)
Louis Verhaardd7911c42020-08-25 13:36:41 +02001027 quantized_min = min(ix)
1028 quantized_max = max(ix)
1029 for x in ix:
1030 if x < zp_in:
1031 lut_result = zp_out + fp_math.multiply_by_quantized_multiplier(
1032 alpha_scalar * (x - zp_in), alpha_scale, alpha_shift
1033 )
1034 else:
1035 lut_result = zp_out + fp_math.multiply_by_quantized_multiplier(x - zp_in, identity_scale, identity_shift)
1036 lut_result = min(quantized_max, max(quantized_min, lut_result))
1037 values.append(lut_result)
Louis Verhaard2e186c72020-10-09 10:47:04 +02001038 return convert_to_lut(op, values, "lrelu")
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001039
1040
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001041def convert_lrelu(op, arch, nng):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001042 # Converts LeakyRelu to a LUT based solution if possible, otherwise a mul + max
Louis Verhaardaee5d752020-09-30 09:01:52 +02001043 if op.type != Op.LeakyRelu:
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001044 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001045 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +01001046 if ifm is None or ofm is None:
1047 return op
Louis Verhaardd7911c42020-08-25 13:36:41 +02001048 if ifm.dtype in (DataType.uint8, DataType.int8) and ifm.dtype == ofm.dtype:
1049 # use LUT for int8/uint8
1050 return convert_lrelu_to_lut(op, arch)
Tim Hall93582962020-09-09 21:58:15 +01001051 if check_quantized_tens_scaling_equal(ifm, ofm) and ifm.dtype == ofm.dtype == DataType.int16:
Louis Verhaardd7911c42020-08-25 13:36:41 +02001052 # use LeakyRelu unmodified for int16 with equal input/output scaling
1053 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001054 return convert_lrelu_to_mul_max(op, arch)
1055
1056
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001057def convert_tanh_sigmoid_to_lut(op, arch, nng):
Louis Verhaardf03bad32020-09-25 08:30:44 +02001058 # Converts int8/uint8 Sigmoid and Tanh to a LUT based solution
Louis Verhaardaee5d752020-09-30 09:01:52 +02001059 if op.type == Op.Sigmoid:
Louis Verhaard2e186c72020-10-09 10:47:04 +02001060 return convert_to_lut8(op, clamp_sigmoid, "sigmoid")
Louis Verhaardaee5d752020-09-30 09:01:52 +02001061 elif op.type == Op.Tanh:
Louis Verhaard2e186c72020-10-09 10:47:04 +02001062 return convert_to_lut8(op, math.tanh, "tanh")
Louis Verhaardf03bad32020-09-25 08:30:44 +02001063 return op
1064
1065
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001066def remove_unwanted_reshapes(op, arch, nng):
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001067 # Try to remove reshapes enclosing ElementWise operator with only one non-constant input
Louis Verhaardaee5d752020-09-30 09:01:52 +02001068 if not op.run_on_npu or not op.type.is_elementwise_op():
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001069 return op
1070
1071 # Check if the ElementWise operator only have one non-constant input
Louis Verhaardaee5d752020-09-30 09:01:52 +02001072 non_const_tens = [x for x in op.inputs if x.ops[0].type != Op.Const]
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001073 if len(non_const_tens) != 1:
1074 return op
1075 ifm = non_const_tens[0]
1076
1077 # Check if operation is enclosed by Reshapes that can be removed
1078 ofm = op.outputs[0]
1079 prev_op = ifm.ops[0]
1080 if (
1081 len(ifm.consumer_list) == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +02001082 and prev_op.type == Op.Reshape
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001083 and len(ofm.consumer_list) == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +02001084 and ofm.consumer_list[0].type == Op.Reshape
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001085 ):
1086 # Operation is enclosed by reshapes, check if they can be removed
Louis Verhaardaee5d752020-09-30 09:01:52 +02001087 prev_op_ifm, prev_op_ofm = prev_op.get_ifm_ofm()
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001088 cons_op = ofm.consumer_list[0]
1089 cons_op_ifm = ofm
1090 cons_op_ofm = cons_op.outputs[0]
1091 if len(prev_op_ifm.shape) == len(cons_op_ofm.shape):
1092 # Check if quantization is the same in the input and output for the reshape ops
Tim Hall93582962020-09-09 21:58:15 +01001093 if check_quantized_tens_scaling_equal(prev_op_ifm, prev_op_ofm) and check_quantized_tens_scaling_equal(
1094 cons_op_ifm, cons_op_ofm
1095 ):
Patrik Gustavsson7ad862a2020-09-29 14:09:43 +02001096 op.set_input_tensor(prev_op_ifm, 0)
1097 op.set_output_tensor(cons_op_ofm)
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001098 return op
1099
1100
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001101def fuse_activation_function_with_prev(op, arch, nng):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001102 # if op is a no-op: attempts to move the activation function to the preceding op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001103 if not op.attrs.get("is_nop", False) or op.activation is None:
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001104 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001105 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +01001106 if ifm is None or ofm is None:
1107 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001108 # finds the input(s) to the operation
1109 prev_op = ifm.ops[0]
1110 # Note: the below checks on prev_op require that a first optimize pass on the full graph has been performed
1111 fuse = (
1112 prev_op.run_on_npu
Louis Verhaardaee5d752020-09-30 09:01:52 +02001113 and prev_op.type.npu_block_type != NpuBlockType.Default
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001114 and len(ifm.ops) == 1
1115 and len(prev_op.outputs[0].consumers()) == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +02001116 and prev_op.activation is None
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001117 )
1118 if op.activation_lut is not None and arch.shram_reserved_unused_banks == 0:
1119 # TODO: if SHRAM LUT space is shared with SHRAM ACC (32, 64 MAC),
1120 # LUT currently only works correctly for elementwise ops
1121 fuse = False
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001122 if not fuse:
1123 return op
1124 # Move the fused activation function + corresponding info to prev_op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001125 prev_op.activation = op.activation
1126 prev_op.forced_output_quantization = op.forced_output_quantization
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001127 if op.activation_lut is not None:
1128 prev_op.set_activation_lut(op.activation_lut)
1129 # Bypass op
Louis Verhaard98a34992020-09-01 10:39:04 +02001130 prev_op.set_output_tensor(ofm)
Tim Halle6ccd872020-11-09 16:46:37 +00001131 DebugDatabase.add_optimised(op, prev_op)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001132 return op
1133
1134
Louis Verhaardae2d5532020-12-11 17:19:54 +01001135def optimise_pad(op, arch, nng):
1136 """
1137 Converts tens1 -> PAD -> tens2 -> CONV to tens1 -> CONV
1138 if both operations can be run on the NPU.
1139 """
1140 if (
1141 (op.type.is_conv2d_op() or op.type.is_depthwise_conv2d_op())
1142 and op.run_on_npu
1143 and op.attrs["padding"] == Padding.VALID
1144 ):
1145 pad_op = op.ifm.ops[0]
1146 if pad_op.type != Op.Pad or not pad_op.run_on_npu:
1147 return op
1148 # Bypass the PAD operator
1149 op.set_input_tensor(pad_op.ifm, 0)
1150 # Adjust the padding attributes of the convolution operator
1151 op.attrs["padding"] = Padding.EXPLICIT
1152 padding = pad_op.inputs[1].values # 4x2 tensor, first dimension is N, H, W, C
1153 top, left, bottom, right = (padding[1][0], padding[2][0], padding[1][1], padding[2][1])
1154 op.attrs["explicit_padding"] = (top, left, bottom, right)
1155 op.set_ifm_ofm_shapes()
1156 return op
1157
1158
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001159def add_attrs_to_resizebilinear(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +02001160 if op.type == Op.ResizeBilinear and op.run_on_npu:
Dwight Lidman42fed942020-05-29 09:37:03 +02001161 input_tensor = op.inputs[0]
1162 upscaled_shape = [input_tensor.shape[1] * 2, input_tensor.shape[2] * 2]
1163 out_shape = op.outputs[0].shape[1:3]
1164 if not op.attrs["align_corners"] and out_shape == upscaled_shape:
1165 # this means the output is supposed to be a x2 upscale,
1166 # so we need to do SAME padding
Michael McGeagh16895482020-12-14 15:51:20 +00001167 op.attrs["padding"] = Padding.SAME
Dwight Lidman42fed942020-05-29 09:37:03 +02001168 elif op.attrs["align_corners"] and out_shape == [upscaled_shape[0] - 1, upscaled_shape[1] - 1]:
1169 # here we can just run the avg pool without padding and
1170 # produce a (M * 2 - 1, N * 2 - 1) sized output
Michael McGeagh16895482020-12-14 15:51:20 +00001171 op.attrs["padding"] = Padding.VALID
Dwight Lidman42fed942020-05-29 09:37:03 +02001172 else:
Charles Xu9a03fdf2020-07-02 15:12:40 +02001173 return op
Dwight Lidman42fed942020-05-29 09:37:03 +02001174 input_tensor.resampling_mode = resampling_mode.NEAREST
Tim Hallc30f4952020-06-15 20:47:35 +01001175 op.attrs.update({"strides": (1, 1, 1, 1), "ksize": (1, 2, 2, 1)})
Dwight Lidman42fed942020-05-29 09:37:03 +02001176 return op
1177
1178
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001179def fixup_bias_tensors(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +02001180 if op.type.needs_bias() and op.bias is None:
Jacob Bohlina41cd4d2020-08-26 18:21:28 +02001181 # Op has no bias, add bias tensor filled with zeros
1182 nr_biases = op.inputs[1].shape[-1]
1183 bias_values = [0] * nr_biases
1184 bias_tensor = create_const_tensor(op.name + "_bias", [nr_biases], DataType.int32, bias_values)
1185 bias_tensor.quant_values = bias_tensor.values
1186 op.set_input_tensor(bias_tensor, -1)
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +02001187
1188 return op
1189
1190
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001191def supported_operator_check(op, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +01001192 op.run_on_npu = arch.supported_operators.is_operator_supported(op)
1193 return op
1194
1195
Tim Halle6ccd872020-11-09 16:46:37 +00001196def _record_optimised(op, arch):
1197 if op.type != Op.Const:
1198 DebugDatabase.add_optimised(op, op)
1199
1200
Tim Hall79d07d22020-04-27 18:20:16 +01001201def optimise_graph_a(nng, arch, verbose_graph=False):
1202 if verbose_graph:
1203 nng.print_graph()
1204
Patrik Gustavsson2349d422020-12-01 16:02:29 +01001205 pre_process_list = [
1206 supported_operator_check,
1207 set_ifm_ofm_op_shapes,
1208 # TODO: memory-only Op removal
1209 ]
1210
1211 for idx, sg in enumerate(nng.subgraphs):
1212 # rewrite graph pass
1213 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
1214 nng, sg, arch, [], pre_process_list, rewrite_unsupported=False,
1215 )
1216
Tim Hall79d07d22020-04-27 18:20:16 +01001217 op_rewrite_list = [
Tim Hall4e127762020-05-15 16:05:49 +01001218 set_tensor_equivalence,
Tim Hall79d07d22020-04-27 18:20:16 +01001219 convert_depthwise_to_conv,
Michael McGeagh8d939c02020-07-29 13:11:43 +01001220 convert_conv_to_fc,
Fredrik Svedberga0c36242020-06-03 15:43:31 +02001221 convert_softmax,
Diqing Zhong016b8272020-12-16 16:46:06 +01001222 optimise_strided_conv,
Tim Hall79d07d22020-04-27 18:20:16 +01001223 fixup_fully_connected_input,
Diqing Zhong94457b12020-12-09 15:22:40 +01001224 convert_batched_fc_shape,
Tim Hall79d07d22020-04-27 18:20:16 +01001225 fixup_pack_input,
Fredrik Svedberg0f98b362020-09-29 10:00:39 +02001226 unfuse_activation_function,
Tim Hall79d07d22020-04-27 18:20:16 +01001227 fixup_conv2d_backprop,
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +01001228 fixup_relus_with_differing_ifm_ofm_scaling,
Tim Hall79d07d22020-04-27 18:20:16 +01001229 fixup_act_reorder,
Charles Xu78792222020-05-13 10:15:26 +02001230 fixup_elementwise_with_scalars,
Jacob Bohline843d332020-06-23 12:12:56 +02001231 reorder_depthwise_weights,
Charles Xu9a03fdf2020-07-02 15:12:40 +02001232 fixup_resizebilinear,
Jacob Bohlina41cd4d2020-08-26 18:21:28 +02001233 fixup_bias_tensors,
Dwight Lidmanc3862c22020-09-14 15:22:33 +02001234 convert_nop_split_to_identity,
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001235 convert_mul_max_to_abs_or_lrelu,
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001236 remove_unwanted_reshapes,
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001237 convert_lrelu,
Louis Verhaardf03bad32020-09-25 08:30:44 +02001238 convert_tanh_sigmoid_to_lut,
Tim Hall79d07d22020-04-27 18:20:16 +01001239 ]
1240
1241 for idx, sg in enumerate(nng.subgraphs):
1242 # rewrite graph pass
1243 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
Dwight Lidman73320a42020-11-05 10:34:41 +01001244 nng, sg, arch, [], op_rewrite_list, rewrite_unsupported=False,
Tim Hall79d07d22020-04-27 18:20:16 +01001245 )
1246
1247 for idx, sg in enumerate(nng.subgraphs):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001248 # remove passthrough tensors and attempt further optimizations
1249 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
Louis Verhaardae2d5532020-12-11 17:19:54 +01001250 nng,
1251 sg,
1252 arch,
1253 [remove_passthrough_tensor],
1254 [fuse_activation_function_with_prev, optimise_pad, add_padding_fields],
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001255 )
Tim Hall79d07d22020-04-27 18:20:16 +01001256
Tim Halle6ccd872020-11-09 16:46:37 +00001257 # Post-optimisation operator debug tracing
1258 for sg in nng.subgraphs:
1259 rewrite_graph.visit_graph_post_order(sg.output_tensors, arch, [], [_record_optimised])
1260
Tim Hall79d07d22020-04-27 18:20:16 +01001261 if verbose_graph:
1262 nng.print_graph()
1263 return nng
1264
Diego Russoea6111a2020-04-14 18:41:58 +01001265
Tim Hall79d07d22020-04-27 18:20:16 +01001266def optimise_graph_b(nng, arch, verbose_graph=False):
1267 if verbose_graph:
1268 nng.print_graph()
1269
1270 for idx, sg in enumerate(nng.subgraphs):
1271 # combined rewrite graph pass
Dwight Lidmanc6ac1942020-10-02 14:55:45 +02001272 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
patrik.gustavssoneeb85152020-12-21 17:10:40 +00001273 nng, sg, arch, [fixup_unpack_output, fixup_stridedslice_output, rewrite_concat, rewrite_split], [],
Dwight Lidmanc6ac1942020-10-02 14:55:45 +02001274 )
Tim Hall79d07d22020-04-27 18:20:16 +01001275
1276 if verbose_graph:
1277 nng.print_graph()
1278 return nng