blob: 3759d3b0d56d5a87f90a9149ec7dd8bc79c8059c [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
159def calc_padding_and_skirt(padding_type, kernel_size, stride, input_dims):
160 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
172 else:
Michael McGeagh16895482020-12-14 15:51:20 +0000173 raise UnsupportedFeatureError(f"Unknown padding")
Tim Hall79d07d22020-04-27 18:20:16 +0100174 padding = (top_pad, left_pad, bottom_pad, right_pad)
175 skirt = (top_pad, left_pad, ypad - top_pad, xpad - left_pad)
176 return padding, skirt
177
Tim Hallc30f4952020-06-15 20:47:35 +0100178
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200179def calc_upscaled_padding_and_skirt(padding_type, kernel_size, stride, input_dims, upscaling_factor):
180 kernel_height, kernel_width = kernel_size[0], kernel_size[1]
Michael McGeagh16895482020-12-14 15:51:20 +0000181 if padding_type == Padding.SAME:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200182 ypad = needed_total_padding(int(input_dims[1]) * upscaling_factor, int(stride[1]), int(kernel_height))
183 xpad = needed_total_padding(int(input_dims[2]) * upscaling_factor, int(stride[2]), int(kernel_width))
Jacob Bohlind47cc272020-08-24 11:42:14 +0200184 right_pad = max(((xpad + 1) // upscaling_factor) - 1, 0)
185 bottom_pad = max(((ypad + 1) // upscaling_factor) - 1, 0)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200186 left_pad = max(kernel_width - 1 - right_pad, 0)
187 top_pad = max(kernel_height - 1 - bottom_pad, 0)
Michael McGeagh16895482020-12-14 15:51:20 +0000188 elif padding_type == Padding.VALID:
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200189 right_pad = max(kernel_width - 2, 0)
190 bottom_pad = max(kernel_height - 2, 0)
191 left_pad = kernel_width - 1
192 top_pad = kernel_height - 1
Jacob Bohlincf7da102020-05-20 09:03:40 +0200193 else:
Michael McGeagh16895482020-12-14 15:51:20 +0000194 raise UnsupportedFeatureError(f"Unknown padding")
Jacob Bohlincf7da102020-05-20 09:03:40 +0200195 padding = (top_pad, left_pad, bottom_pad, right_pad)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200196 skirt = padding
Jacob Bohlincf7da102020-05-20 09:03:40 +0200197 return padding, skirt
198
Tim Hall79d07d22020-04-27 18:20:16 +0100199
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200200def fixup_conv2d_backprop(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200201 if op.type == Op.Conv2DBackpropInput:
Tim Hall79d07d22020-04-27 18:20:16 +0100202 # flip the inputs
203 op.inputs[0], op.inputs[2] = op.inputs[2], op.inputs[0]
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000204 op.set_ifm_ofm_shapes()
Louis Verhaardaee5d752020-09-30 09:01:52 +0200205 op.type = Op.Conv2DBackpropInputSwitchedBias
Jacob Bohlincf7da102020-05-20 09:03:40 +0200206
207 # Update strides
Tim Hallc30f4952020-06-15 20:47:35 +0100208 op.attrs.update({"stride_w": 1, "stride_h": 1, "strides": (1, 1, 1, 1)})
Tim Hall79d07d22020-04-27 18:20:16 +0100209
210 return op
211
212
Charles Xu9a03fdf2020-07-02 15:12:40 +0200213# Convert the op to an elementwise add
214def convert_resizebilinear_1x1_to_add(op):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200215 op.type = Op.Add
Charles Xu9a03fdf2020-07-02 15:12:40 +0200216 op.name = op.name + "_add"
Charles Xu9a03fdf2020-07-02 15:12:40 +0200217 op.attrs["resizebilinear"] = True
218 # Create an input tensor filled with zeros
219 shape = op.outputs[0].shape
220 tens = Tensor(shape, op.inputs[0].dtype, op.inputs[1].name + "_add")
221 tens.values = np.zeros(shape)
222 tens.quant_values = np.zeros(shape, np.uint8)
223 tens.quantization = QuantizationParameters(0.0, 255.0)
224 tens.quantization.scale_f32 = 1.0
225 tens.quantization.zero_point = 0
226 tens.consumer_list = [op]
227 tens_op = op.inputs[1].ops[0]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100228 tens_op.set_output_tensor(tens)
Charles Xu9a03fdf2020-07-02 15:12:40 +0200229 # Set the add inputs
230 op.inputs[1] = op.inputs[0]
231 op.inputs[0] = tens
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000232 op.set_ifm_ofm_shapes()
Charles Xu9a03fdf2020-07-02 15:12:40 +0200233
234 return op
235
236
Charles Xu87c13502020-08-06 12:17:26 +0200237# Convert ResizeBilinear to a number of 2x2 pool ops
238def convert_resizebilinear_to_2x2_pool(op):
239 count = 0
240 pre_op = op
241 outputs = op.outputs
242
243 op.attrs.update({"strides": (1, 1, 1, 1), "ksize": (1, 2, 2, 1)})
244 if op.attrs["align_corners"]:
245 shape_modifier = 1
Michael McGeagh16895482020-12-14 15:51:20 +0000246 op.attrs["padding"] = Padding.VALID
Charles Xu87c13502020-08-06 12:17:26 +0200247 else:
248 shape_modifier = 0
Michael McGeagh16895482020-12-14 15:51:20 +0000249 op.attrs["padding"] = Padding.SAME
Charles Xu87c13502020-08-06 12:17:26 +0200250 op.inputs[0].resampling_mode = resampling_mode.NEAREST
251
252 upscaled_shape = np.array(op.inputs[0].shape[1:3])
253 out_shape = np.array(op.outputs[0].shape[1:3])
254 if (upscaled_shape == upscaled_shape * 2 - shape_modifier).all():
255 return op
256
257 while (upscaled_shape < out_shape).all():
258 if count == 0:
259 scaled_op = pre_op
260 else:
261 scaled_op = op.clone("_{}".format(count))
262 scaled_op.inputs[0] = pre_op.outputs[0]
263
264 upscaled_shape = upscaled_shape * 2 - shape_modifier
265
266 if (upscaled_shape == out_shape).all():
267 scaled_op.outputs = outputs
268 scaled_op.outputs[0].ops = [scaled_op]
269 else:
270 shape = outputs[0].shape.copy()
271 shape[1:3] = upscaled_shape[0:2]
272 out_tens = Tensor(shape, DataType.int16, "{}_{}".format(op.outputs[0].name, count))
273 out_tens.quantization = op.outputs[0].quantization.clone()
274 out_tens.quantization.quant_min = np.iinfo(np.int16).min
275 out_tens.quantization.quant_max = np.iinfo(np.int16).max
276 scaled_op.set_output_tensor(out_tens)
277 pre_op = scaled_op
278 count += 1
279
280 # Setup the scale value
281 if scaled_op.inputs[0].dtype.bits == 8 and scaled_op.outputs[0].dtype.bits == 16:
282 scaled_op.attrs["rescale"] = 128
283 elif scaled_op.inputs[0].dtype.bits == 16 and scaled_op.outputs[0].dtype.bits == 8:
284 scaled_op.attrs["rescale"] = 1 / 128
285 elif "rescale" in scaled_op.attrs:
286 del scaled_op.attrs["rescale"]
Patrik Gustavssoncc6915c2020-12-22 09:16:50 +0100287 scaled_op.set_ifm_ofm_shapes()
Charles Xu87c13502020-08-06 12:17:26 +0200288
289 return op
290
291
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200292def fixup_resizebilinear(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200293 if op.type == Op.ResizeBilinear and op.run_on_npu:
Charles Xu87c13502020-08-06 12:17:26 +0200294 if op.inputs[0].shape == op.outputs[0].shape:
Charles Xu36ffaf32020-08-05 15:40:44 +0200295 # Bypass nop resizebilinear
296 op.inputs = op.inputs[:1]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200297 op.type = Op.Identity
Charles Xu87c13502020-08-06 12:17:26 +0200298 elif op.inputs[0].shape[1] == 1 and op.inputs[0].shape[2] == 1:
299 convert_resizebilinear_1x1_to_add(op)
300 else:
301 convert_resizebilinear_to_2x2_pool(op)
Charles Xu9a03fdf2020-07-02 15:12:40 +0200302
303 return op
304
305
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200306def convert_nop_split_to_identity(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200307 if op.type == Op.Split and op.attrs.get("num_splits") == 1:
Dwight Lidmanc3862c22020-09-14 15:22:33 +0200308 # the list comprehension should return a list with a single tensor
309 # if it shouldn't, remove_passthrough_tensor will fail appropriately
310 op.inputs = [i for i in op.inputs if i.shape == op.outputs[0].shape]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200311 op.type = Op.Identity
Dwight Lidmanc3862c22020-09-14 15:22:33 +0200312 return op
313
314
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200315def fixup_fully_connected_input(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200316 if op.type == Op.FullyConnected:
Tim Hall79d07d22020-04-27 18:20:16 +0100317 inp = op.inputs[0]
318 weights = op.inputs[1]
319
320 n_in_elems = weights.shape[-2]
321 elms = inp.elements()
322 batch_size = elms // n_in_elems
323 assert batch_size * n_in_elems == elms
324
325 desired_shape = [batch_size, n_in_elems]
326 if inp.shape != desired_shape:
327 # mismatch, insert a reshape to fix this.
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200328 op.set_input_tensor(create_reshape_tensor(inp, desired_shape), 0)
Tim Hall79d07d22020-04-27 18:20:16 +0100329
330 return op
331
332
Diqing Zhong94457b12020-12-09 15:22:40 +0100333def convert_batched_fc_shape(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200334 if op.type == Op.FullyConnected:
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200335 ifm = op.inputs[0]
336 ofm = op.outputs[0]
337 # Check if the FC is 2D and first dimension indicates batching
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100338 # TOD0 op.ifm_shape[0] > 1 is enough when refactory is complete
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000339 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 +0200340 n = ifm.shape[0]
341 batching_split = {4: (2, 2), 8: (2, 4), 16: (4, 4)}
342 h, w = batching_split.get(n, (1, n))
343
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200344 prev_op = ifm.ops[0]
345 desired_shape = [1, h, w, ifm.shape[-1]]
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000346 op.ifm_shapes[0] = Shape4D(desired_shape)
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100347
Louis Verhaardaee5d752020-09-30 09:01:52 +0200348 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 +0200349 # There is a preceding Reshape
350 # Compare input of prev_op and input of op, to see if prev_op can be removed
351 ifm_prev_op = prev_op.inputs[0]
Tim Hall89567612020-10-27 11:57:57 +0000352 if ifm_prev_op.shape == ifm.shape and check_quantized_tens_scaling_equal(ifm_prev_op, ifm):
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200353 # prev_op can be removed
354 op.set_input_tensor(ifm_prev_op, 0)
355 else:
356 op.inputs[0].set_all_shapes(desired_shape)
357 prev_op.set_input_tensor(
358 create_const_tensor(prev_op.inputs[1].name, [1], DataType.int32, desired_shape), 1
359 )
360 prev_op.attrs["new_shape"] = desired_shape
361 else:
362 # Add reshape op to the input if there is no preceding reshape
363 ifm.consumer_list.remove(op)
364 op.set_input_tensor(create_reshape_tensor(ifm, desired_shape), 0)
365
366 # Reshape Weights to be 4D. IO becomes HWIO
367 weight_tensor = op.inputs[1]
368 weight_tensor.quant_values = np.expand_dims(np.expand_dims(weight_tensor.quant_values, axis=0), axis=0)
369 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
370
371 desired_shape = [1, h, w, ofm.shape[-1]]
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000372 op.ofm_shapes[0] = Shape4D(desired_shape)
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100373
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200374 if (
375 len(ofm.consumer_list) == 1
376 and ofm.consumer_list[0] is not None
Louis Verhaardaee5d752020-09-30 09:01:52 +0200377 and ofm.consumer_list[0].type == Op.Reshape
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200378 ):
379 # There is a subsequent Reshape
380 # Compare desired shape and output of consumer op, to see if consumer op can be removed
381 ofm_cons_op = ofm.consumer_list[0].outputs[0]
Tim Hall93582962020-09-09 21:58:15 +0100382 if desired_shape == ofm_cons_op.shape and check_quantized_tens_scaling_equal(ofm, ofm_cons_op):
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200383 op.outputs[0] = ofm_cons_op
384 op.outputs[0].ops = [op]
385 else:
386 op.outputs[0].set_all_shapes(desired_shape)
387 else:
Diqing Zhong94457b12020-12-09 15:22:40 +0100388 # Add reshape op to the output
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200389 op.set_output_tensor(create_reshape_tensor(ofm, desired_shape, False))
390 return op
391
392
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200393def fixup_pack_input(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200394 if op.type == Op.Pack:
Tim Hall79d07d22020-04-27 18:20:16 +0100395 # Pack is also referred to as Stack
396 # Requires the rewrite_concat function to be called on the op afterwards
397 axis = int(op.attrs["axis"])
398 desired_shape = op.inputs[0].shape[:axis] + [1] + op.inputs[0].shape[axis:]
399
400 # Construct 1 shape tensor to be used by all inserted reshape ops
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100401 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, desired_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100402
403 for idx, inp in enumerate(op.inputs):
Tim Hall79d07d22020-04-27 18:20:16 +0100404 reshape_out = inp.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100405 reshape_out.set_all_shapes(desired_shape)
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100406
Louis Verhaardaee5d752020-09-30 09:01:52 +0200407 reshape_op = Operation(Op.Reshape, "{}{}_reshape".format(op.name, idx))
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100408 reshape_op.attrs["new_shape"] = desired_shape
409 reshape_op.inputs = [inp, new_shape_tens]
410 reshape_op.set_output_tensor(reshape_out)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000411 reshape_op.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000412 DebugDatabase.add_optimised(op, reshape_op)
Tim Hall79d07d22020-04-27 18:20:16 +0100413
414 op.inputs[idx] = reshape_out
415
Louis Verhaardaee5d752020-09-30 09:01:52 +0200416 op.type = Op.PackReshaped
Tim Hall79d07d22020-04-27 18:20:16 +0100417
418 return op
419
420
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200421def unfuse_activation_function(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200422 if op.type == Op.ConcatTFLite and op.run_on_npu and op.activation is not None:
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100423 act_op = Operation(op.activation.op_type, op.name + op.activation.op_type.name)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200424 op.activation = None
Fredrik Svedberg0f98b362020-09-29 10:00:39 +0200425 out_tens = op.outputs[0]
426 intermediate_tens = out_tens.clone("_act_intermediate")
427 act_op.set_output_tensor(out_tens)
428 act_op.add_input_tensor(intermediate_tens)
429 op.set_output_tensor(intermediate_tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000430 act_op.set_ifm_ofm_shapes()
Fredrik Svedberg0f98b362020-09-29 10:00:39 +0200431
432 return op
433
Louis Verhaard8912c532020-09-30 12:11:49 +0200434
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100435def fixup_stridedslice_output(tens, arch, nng):
436 op = tens.ops[0]
Dwight Lidman73320a42020-11-05 10:34:41 +0100437 if op.run_on_npu and op.type == Op.StridedSlice:
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100438 reshape_input_shape = tens.shape
439 new_axis_mask = op.attrs["new_axis_mask"]
440 shrink_axis_mask = op.attrs["shrink_axis_mask"]
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100441
Dwight Lidman73320a42020-11-05 10:34:41 +0100442 if shrink_axis_mask != 0:
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100443 n = 0
444 axis = 0
445 while shrink_axis_mask:
446 prev_mask = shrink_axis_mask
447 n += 1
448 shrink_axis_mask &= shrink_axis_mask - 1
449 axis = int(math.log2(prev_mask - shrink_axis_mask))
450 reshape_input_shape = reshape_input_shape[:axis] + [1] + reshape_input_shape[axis:]
451
452 assert len(tens.shape) == (len(op.inputs[0].shape) - n)
453 op.attrs["shrink_axis_mask"] = 0
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100454 elif new_axis_mask != 0:
455 n = 0
456 axis = 0
457 while new_axis_mask:
458 prev_mask = new_axis_mask
459 n += 1
460 new_axis_mask &= new_axis_mask - 1
461 axis = int(math.log2(prev_mask - new_axis_mask))
462 reshape_input_shape = reshape_input_shape[:axis] + reshape_input_shape[(axis + 1) :]
463 new_axis_mask >>= 1
464
465 assert len(tens.shape) == (len(op.inputs[0].shape) + n)
466 op.attrs["new_axis_mask"] = 0
Dwight Lidmandda21af2020-11-11 15:44:57 +0100467 else:
468 # Equal Rank StridedSlice, no need to insert reshape
469 return tens
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100470
471 # Construct 1 shape tensor to be used by all inserted reshape ops
472 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, tens.shape)
473
474 for idx, out_tens in enumerate(op.outputs):
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000475 op.ofm_shapes[idx] = Shape4D(new_shape_tens.shape)
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100476 reshape_in = out_tens.clone("_reshaped")
477 reshape_in.set_all_shapes(reshape_input_shape)
478 reshape_in.ops = [op]
479
480 reshape_op = Operation(Op.Reshape, "{}{}_reshape".format(op.name, idx))
481 reshape_op.attrs["new_shape"] = reshape_input_shape
482 reshape_op.inputs = [reshape_in, new_shape_tens]
483 reshape_op.set_output_tensor(out_tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000484 reshape_op.set_ifm_ofm_shapes()
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100485
486 op.outputs[idx] = reshape_in
487
488 return tens
489
490
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200491def fixup_unpack_output(tens, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +0100492 op = tens.ops[0]
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100493 if op.run_on_npu and op.type == Op.Unpack:
Tim Hall79d07d22020-04-27 18:20:16 +0100494 # Unpack is also referred to as Unstack
495 # Requires the rewrite_split function to be called on the op afterwards
Diqing Zhongc7c0b1b2020-10-26 11:45:25 +0100496 axis = int(op.attrs["axis"])
497 op.type = Op.UnpackReshaped
498 reshape_input_shape = tens.shape[:axis] + [1] + tens.shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100499
500 # Construct 1 shape tensor to be used by all inserted reshape ops
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100501 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, tens.shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100502
503 for idx, out_tens in enumerate(op.outputs):
Tim Hall79d07d22020-04-27 18:20:16 +0100504 reshape_in = out_tens.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100505 reshape_in.set_all_shapes(reshape_input_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100506 reshape_in.ops = [op]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100507
Louis Verhaardaee5d752020-09-30 09:01:52 +0200508 reshape_op = Operation(Op.Reshape, "{}{}_reshape".format(op.name, idx))
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100509 reshape_op.attrs["new_shape"] = reshape_input_shape
Tim Hall79d07d22020-04-27 18:20:16 +0100510 reshape_op.inputs = [reshape_in, new_shape_tens]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100511 reshape_op.set_output_tensor(out_tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000512 reshape_op.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000513 DebugDatabase.add_optimised(op, reshape_op)
Tim Hall79d07d22020-04-27 18:20:16 +0100514
515 op.outputs[idx] = reshape_in
Tim Hall79d07d22020-04-27 18:20:16 +0100516 return tens
517
518
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200519def add_padding_fields(op, arch, nng):
Jacob Bohlin90033f32020-08-28 15:45:44 +0200520 if op.run_on_npu:
521 if "padding" in op.attrs:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200522 if op.type.is_conv2d_op() or op.type.is_depthwise_conv2d_op():
Jacob Bohlin90033f32020-08-28 15:45:44 +0200523 kernel_size = op.inputs[1].shape[:2]
524 input_shape = op.inputs[0].shape
Louis Verhaardaee5d752020-09-30 09:01:52 +0200525 elif op.type.is_pool_op() or op.type.npu_block_type == NpuBlockType.ReduceSum:
Jacob Bohlin90033f32020-08-28 15:45:44 +0200526 kernel_size = op.attrs["ksize"][1:3]
527 input_shape = op.inputs[0].shape
Jacob Bohlin90033f32020-08-28 15:45:44 +0200528 else:
Michael McGeagh7a6f8432020-12-02 15:29:22 +0000529 raise UnsupportedFeatureError(f"Unknown operation that uses padding: {optype_to_builtintype(op.type)}")
Tim Hall79d07d22020-04-27 18:20:16 +0100530
Louis Verhaardaee5d752020-09-30 09:01:52 +0200531 if op.type == Op.Conv2DBackpropInputSwitchedBias:
Jacob Bohlin90033f32020-08-28 15:45:44 +0200532 upscaling_factor = op.outputs[0].shape[1] // input_shape[1]
533 padding, skirt = calc_upscaled_padding_and_skirt(
534 op.attrs["padding"], kernel_size, op.attrs["strides"], input_shape, upscaling_factor
535 )
536 else:
537 dilation_h, dilation_w = op.get_dilation_h_w()
538 dilated_kernel_size = [dilation_h * (kernel_size[0] - 1) + 1, dilation_w * (kernel_size[1] - 1) + 1]
539 padding, skirt = calc_padding_and_skirt(
540 op.attrs["padding"], dilated_kernel_size, op.attrs["strides"], input_shape
541 )
Jacob Bohlincf7da102020-05-20 09:03:40 +0200542
Jacob Bohlin90033f32020-08-28 15:45:44 +0200543 op.attrs["explicit_padding"] = padding
544 op.attrs["skirt"] = skirt
Jacob Bohlincf7da102020-05-20 09:03:40 +0200545
Tim Hall79d07d22020-04-27 18:20:16 +0100546 return op
547
548
Tim Hall79d07d22020-04-27 18:20:16 +0100549# Check if the op can be reordered
550def get_prepend_op(op):
551 inp = op.inputs[0]
552 # The op should be reordered between prev_op and prep_op
553 prev_op = inp.ops[-1]
554 prep_op = None
555 while prev_op.type in memory_only_ops and len(prev_op.outputs) == 1 and len(prev_op.outputs[0].consumers()) == 1:
556 prep_op = prev_op
557 inp = prev_op.inputs[0]
558 prev_op = inp.ops[-1]
Diego Russoea6111a2020-04-14 18:41:58 +0100559 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 +0100560 return prep_op
561
562 return None
563
564
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200565def convert_depthwise_to_conv(op, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +0100566 # Depthwise is equivalent to a single conv2d if the ifm depth is 1 and
567 # the ofm depth equals the depth multipler.
568 # If those conditions are true, then we can perform a simple
569 # switch of the operator type (and weight order)
570
Louis Verhaardaee5d752020-09-30 09:01:52 +0200571 if op.type == Op.DepthwiseConv2DBias and (op.attrs["depth_multiplier"] != 1):
Tim Hall79d07d22020-04-27 18:20:16 +0100572 ifm_tensor = op.inputs[0]
573 weight_tensor = op.inputs[1]
574 ofm_tensor = op.outputs[0]
575 if (ifm_tensor.shape[3] == 1) and (ofm_tensor.shape[3] == op.attrs["depth_multiplier"]):
576 # Change op type to Conv2d
Louis Verhaardaee5d752020-09-30 09:01:52 +0200577 op.type = Op.Conv2DBias
Tim Hall79d07d22020-04-27 18:20:16 +0100578 del op.attrs["channel_multiplier"]
579 del op.attrs["depth_multiplier"]
580
581 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100582 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Tim Hall79d07d22020-04-27 18:20:16 +0100583 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200584 raise UnsupportedFeatureError(
Michael McGeagh7a6f8432020-12-02 15:29:22 +0000585 f"Unsupported 'DEPTHWISE_CONV_2D' with depth_multiplier = {op.attrs['depth_multiplier']},",
586 f" ifm channels = {ifm_tensor.shape[3]}, ofm channels = {ofm_tensor.shape[3]}",
Tim Hall79d07d22020-04-27 18:20:16 +0100587 )
Tim Halle6ccd872020-11-09 16:46:37 +0000588 DebugDatabase.add_optimised(op, op)
Tim Hall79d07d22020-04-27 18:20:16 +0100589 return op
590
591
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200592def reorder_depthwise_weights(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200593 if op.type.is_depthwise_conv2d_op():
Jacob Bohline843d332020-06-23 12:12:56 +0200594 weight_tensor = op.inputs[1]
595 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100596 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Jacob Bohline843d332020-06-23 12:12:56 +0200597 weight_tensor.weight_transpose_depthwise = True
598
599 return op
600
601
Diqing Zhong016b8272020-12-16 16:46:06 +0100602def optimise_strided_conv(op, arch, nng):
603 stride_x, stride_y = op.get_kernel_stride()
604 ifm_tensor, _, weight_tensor, _ = op.get_ifm_ifm2_weights_ofm()
605
606 if (
607 op.type == Op.Conv2DBias
608 and op.op_index == 0
609 and stride_x == 2
610 and len(ifm_tensor.shape) == 4
611 and ifm_tensor.shape[3] <= 4
612 and ifm_tensor.shape[2] % 2 == 0
613 and weight_tensor is not None
614 and weight_tensor.shape[1] >= 2
615 ):
616 # IFM
617 ifm_reshaped = create_reshape_tensor(
618 ifm_tensor, [ifm_tensor.shape[0], ifm_tensor.shape[1], ifm_tensor.shape[2] // 2, ifm_tensor.shape[3] * 2]
619 )
620 op.set_input_tensor(ifm_reshaped, 0)
621
622 # Weights
623 weight_shape = weight_tensor.shape
624 if weight_shape[1] % 2 != 0:
625 weight_shape[1] = weight_shape[1] + 1
626 padded_array = np.zeros(weight_shape)
627 for i in range(weight_shape[0]):
628 padded_array[i] = np.vstack(
629 [
630 weight_tensor.quant_values[i],
631 np.full((1, weight_shape[2], weight_shape[3]), weight_tensor.quantization.zero_point),
632 ]
633 )
634 weight_tensor.quant_values = padded_array
635 weight_shape[1] //= 2
636 weight_shape[2] *= 2
637 weight_tensor.quant_values = np.reshape(weight_tensor.quant_values, weight_shape)
638 weight_tensor.set_all_shapes(weight_shape)
639 # If multiple copies of the weights are used, we could avoid
640 # them having the same address by changing the value_id
641 weight_tensor.value_id = uuid.uuid4()
642
643 # Strides
644 stride_x = 1
645 op.attrs.update({"stride_w": stride_x, "stride_h": stride_y, "strides": (1, stride_y, stride_x, 1)})
646
647 op.set_ifm_ofm_shapes()
648
649 return op
650
651
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200652def convert_conv_to_fc(op, arch, nng):
Michael McGeagh8d939c02020-07-29 13:11:43 +0100653 # Conv 1x1 can be equivalent to Fully Connected.
654 # By representing certain convs as fully connected layers, Vela can better determine wether or not to use
655 # caching/double buffering for the weights.
656 # (Weights dont need to be reloaded for convs when IFM H and W are 1)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200657 if op.type == Op.Conv2DBias:
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000658 h = op.ifm_shapes[0].height
659 w = op.ifm_shapes[0].width
Michael McGeagh8d939c02020-07-29 13:11:43 +0100660 kh, kw, _, _ = op.inputs[1].shape
661 if h == 1 and w == 1 and kh == 1 and kw == 1:
662 # Overwrite this op as a Fully Connected Op
663 op.name += "_fc"
Louis Verhaardaee5d752020-09-30 09:01:52 +0200664 op.type = Op.FullyConnected
Michael McGeagh8d939c02020-07-29 13:11:43 +0100665 op.attrs = {
Michael McGeagh8d939c02020-07-29 13:11:43 +0100666 "weights_format": 0,
Michael McGeagh8d939c02020-07-29 13:11:43 +0100667 }
668 # Reshape Weights to be 2D. HWIO becomes just IO (as H and W are 1, they can just be dropped)
669 weight_tensor = op.inputs[1]
670 weight_tensor.quant_values = weight_tensor.quant_values.squeeze(axis=(0, 1))
671 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100672
Michael McGeagh8d939c02020-07-29 13:11:43 +0100673 # The output from a fully connected is expected to be 2D so we need to add a reshape layer to convert it
674 # back to 4D afterwards as the next layer is expecting that shape
675 orig_ofm_tensor = op.outputs[0]
676 # 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})
677 fc_ofm_tensor = orig_ofm_tensor.clone("_fc")
678 fc_ofm_tensor.set_all_shapes([1, fc_ofm_tensor.shape[-1]])
679 fc_ofm_tensor.ops = [op]
680 # Add a reshape after the new OFM to convert it back to the original 4D shape
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100681 reshape_name = op.name + "_reshape"
682 new_shape_tens = create_const_tensor(reshape_name + "_shape", [1], DataType.int32, orig_ofm_tensor.shape)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200683 reshape_op = Operation(Op.Reshape, reshape_name)
Michael McGeagh8d939c02020-07-29 13:11:43 +0100684 reshape_op.attrs["new_shape"] = orig_ofm_tensor.shape
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100685 reshape_op.inputs = [fc_ofm_tensor, new_shape_tens]
686 reshape_op.set_output_tensor(orig_ofm_tensor)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000687 reshape_op.set_ifm_ofm_shapes()
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100688
Michael McGeagh8d939c02020-07-29 13:11:43 +0100689 # Replace this ops OFM to point to the 2D tensor
690 op.outputs[0] = fc_ofm_tensor
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000691 op.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000692 # Record optimisation in debug database
693 DebugDatabase.add_optimised(op, reshape_op)
694 DebugDatabase.add_optimised(op, op)
Michael McGeagh8d939c02020-07-29 13:11:43 +0100695 return op
696
697
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200698def fixup_relus_with_differing_ifm_ofm_scaling(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200699 if op.run_on_npu and op.type.is_relu_op():
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100700 ifm = op.inputs[0]
701 ofm = op.outputs[0]
702 # Relu with differing IFM and OFM scaling cannot be fused with another primary op
703 # and requires its own to be inserted
Tim Hall93582962020-09-09 21:58:15 +0100704 if not check_quantized_tens_scaling_equal(ifm, ofm):
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100705 # Override this op with its own primary op (avgpool)
706 relu_fused_op = create_avgpool_nop(op.name + "_avgpool")
707 # And fuse the original activation function to it
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100708 relu_fused_op.activation = create_activation_function(op.type)
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100709 # Tidy up and assign the ifm and ofm to the new op
710 ifm.consumer_list.remove(op)
Andreas Nevalainenf3d737e2020-09-25 14:12:43 +0200711
712 # if not 4d, reshape ifm/ofm
713 if len(ifm.shape) < 4:
714 ifm_shaped = create_reshape_tensor(ifm, full_shape(4, ifm.shape, 1))
715 ifm = ifm_shaped
716 if len(ofm.shape) < 4:
717 ofm_shaped = create_reshape_tensor(ofm, full_shape(4, ofm.shape, 1), False)
718 ofm = ofm_shaped
719
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100720 relu_fused_op.add_input_tensor(ifm)
721 relu_fused_op.set_output_tensor(ofm)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000722 relu_fused_op.set_ifm_ofm_shapes()
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +0100723 op = relu_fused_op
724 return op
725
726
Tim Hall79d07d22020-04-27 18:20:16 +0100727# Reorder activation op if it's after the memory only operations
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200728def fixup_act_reorder(op, arch, nng):
Michael McGeaghf3e3ad72020-12-02 12:39:03 +0000729 if op.type.is_relu_op() or op.type in (Op.Sigmoid, Op.Tanh):
Tim Hall79d07d22020-04-27 18:20:16 +0100730 prep_op = get_prepend_op(op)
Diego Russoea6111a2020-04-14 18:41:58 +0100731 if prep_op is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100732 act_op = op.clone("_reordered")
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100733 act_op.ifm_shapes = list(op.ifm_shapes)
734 act_op.ofm_shapes = list(op.ofm_shapes)
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200735
736 # There is only one input tensor, overwrite it
737 act_op.set_input_tensor(prep_op.inputs[0], 0)
738
Tim Hall79d07d22020-04-27 18:20:16 +0100739 act_op_out = act_op.inputs[0].clone("_acted")
740 act_op_out.quantization = op.outputs[0].quantization.clone()
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100741 act_op.set_output_tensor(act_op_out)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000742 act_op.ifm_shapes[0] = Shape4D(prep_op.inputs[0].shape)
743 act_op.ofm_shapes[0] = Shape4D(act_op_out.shape)
Patrik Gustavssoncb337042020-09-16 14:55:40 +0200744
745 # Update the consumer list
746 act_op_out.consumer_list = op.outputs[0].consumer_list.copy()
747 act_op_out.consumer_list.append(prep_op)
748
Tim Hall79d07d22020-04-27 18:20:16 +0100749 prep_op.inputs[0] = act_op_out
750 prep_op.outputs[0].quantization = act_op_out.quantization.clone()
751
752 # Mark the op so that it will be removed as passthrough later on
Louis Verhaardaee5d752020-09-30 09:01:52 +0200753 op.type = Op.Identity
Tim Halle6ccd872020-11-09 16:46:37 +0000754
755 # Record optimisation in debug database
756 DebugDatabase.add_optimised(op, act_op)
757 DebugDatabase.add_optimised(op, op)
Tim Hall79d07d22020-04-27 18:20:16 +0100758 return op
759
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200760
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200761def fixup_elementwise_with_scalars(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200762 if op.type.is_binary_elementwise_op():
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200763 ifm_tensor, ifm2_tensor, _, _ = op.get_ifm_ifm2_weights_ofm()
Charles Xu78792222020-05-13 10:15:26 +0200764 if ifm2_tensor.shape != [] and ifm_tensor.shape != []:
765 diff = len(ifm_tensor.shape) - len(ifm2_tensor.shape)
766 if diff > 0:
767 ifm2_tensor.shape = full_shape(len(ifm_tensor.shape), ifm2_tensor.shape, 1)
768 elif diff < 0:
769 ifm_tensor.shape = full_shape(len(ifm2_tensor.shape), ifm_tensor.shape, 1)
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200770 elif ifm_tensor.shape == [] and ifm_tensor.quant_values is None:
771 # IFM is marked as a scalar, but is a result of an operation; change it to a shape of size 1
772 ifm_tensor.shape = len(ifm2_tensor.shape) * [1]
773 ifm_tensor.storage_shape = ifm_tensor.shape
774 elif ifm2_tensor.shape == [] and ifm2_tensor.quant_values is None:
775 # IFM2 is marked as a scalar, but is a result of an operation; change it to a shape of size 1
776 ifm2_tensor.shape = len(ifm_tensor.shape) * [1]
777 ifm2_tensor.storage_shape = ifm2_tensor.shape
Charles Xu78792222020-05-13 10:15:26 +0200778 return op
Tim Hall79d07d22020-04-27 18:20:16 +0100779
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200780
Tim Hall4e127762020-05-15 16:05:49 +0100781# Set input/output tensor equivalence to the same id for memory operations
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200782def set_tensor_equivalence(op, arch, nng):
Michael McGeagh11b0bdb2020-09-08 11:07:35 +0100783 if op.type in memory_only_ops:
Tim Hall4e127762020-05-15 16:05:49 +0100784 eid = op.outputs[0].equivalence_id
785 for inp in op.inputs:
786 inp.equivalence_id = eid
787 return op
788
789
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100790def set_ifm_ofm_op_shapes(op, arch, nng):
791 if op.run_on_npu and op.type.needs_shapes():
792 if op.ifm_shapes or op.ofm_shapes:
793 # Shapes already set
794 return op
795 op.set_ifm_ofm_shapes()
796 return op
797
798
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200799def convert_softmax(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200800 if op.type == Op.Softmax and op.run_on_npu:
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200801 softmax = SoftMax(op)
802 op = softmax.get_graph()
803 return op
804
805
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +0200806def convert_mul_max_to_abs_or_lrelu(op, arch, nng):
Diego Russoea6111a2020-04-14 18:41:58 +0100807 r"""Whenever there is a subgraph with this topology:
Tim Hall79d07d22020-04-27 18:20:16 +0100808
809 Input X For X = -1 or X > 0
810 | \ / This subgraph can be replaced with either
811 | Mul an Abs (if X = -1) or a LeakyReLU (if X > 0)
812 | /
813 Max
814 """
815
Louis Verhaardaee5d752020-09-30 09:01:52 +0200816 if op.type == Op.Maximum:
Tim Hall79d07d22020-04-27 18:20:16 +0100817 # finds the Mul input(s) to the Max
Louis Verhaardaee5d752020-09-30 09:01:52 +0200818 muls = [i for i in op.inputs if i.ops[0].type == Op.Mul]
Tim Hall79d07d22020-04-27 18:20:16 +0100819 if len(muls) == 1:
820 mul = muls[0].ops[0]
821 elif len(muls) == 2:
822 # In the case both inputs are Muls, find the one with the same input as the Max
823 mul = [m for m in muls if len(set(op.inputs + m.ops[0].inputs)) == 1][0].ops[0]
824 else:
825 # No Mul inputs
826 return op
827
828 # make sure the Mul doesn't have any other consumers
Louis Verhaardd7911c42020-08-25 13:36:41 +0200829 mul_ofm = mul.outputs[0]
830 if len(mul_ofm.consumers()) != 1:
Tim Hall79d07d22020-04-27 18:20:16 +0100831 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +0200832 # make sure the Mul doesn't have a fused activation function
833 if mul.activation:
Tim Hall79d07d22020-04-27 18:20:16 +0100834 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +0200835 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +0100836 if ifm is None or ofm is None:
837 return op
838
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200839 if ifm.dtype not in (DataType.uint8, DataType.int8) or ifm.dtype != ofm.dtype:
840 return op
Tim Hall93582962020-09-09 21:58:15 +0100841 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 +0200842 # rewrite to LeakyRelu currently only makes sense if the quantization is identical
843 return op
Tim Hall79d07d22020-04-27 18:20:16 +0100844
845 # finds the branched input that goes to both the Max and the Mul
846 shared = set(op.inputs) & set(mul.inputs)
847 if len(shared) == 1:
848 shared_in = shared.pop()
849 # find the constant scalar input to the Mul
850 const_tens = (set(mul.inputs) - {shared_in}).pop()
851 # check that it is a scalar
852 if const_tens.shape != []:
853 return op
854 const = const_tens.ops[0]
855 # check that it is a constant
Louis Verhaardaee5d752020-09-30 09:01:52 +0200856 if const.type != Op.Const:
Tim Hall79d07d22020-04-27 18:20:16 +0100857 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200858 # Remove the Mul from the shared input's consumers
859 shared_in.consumer_list.remove(mul)
Tim Hall79d07d22020-04-27 18:20:16 +0100860 else:
861 return op
862
863 val = const.outputs[0].values
864 if val >= 0:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200865 new_op = Op.LeakyRelu
Tim Hall79d07d22020-04-27 18:20:16 +0100866 op.attrs["alpha"] = val
Louis Verhaardd7911c42020-08-25 13:36:41 +0200867 # to produce bit exact results, the alpha is not enough;
868 # save additional scaling info in attr "alpha_scale", to be used as input
869 # to the LUT construction
870 alpha_scalar = const_tens.quant_values - const_tens.quantization.zero_point
871 mul_ifm_scale = np.double(ifm.quantization.scale_f32)
872 mul_ifm2_scale = np.double(const_tens.quantization.scale_f32)
873 mul_ofm_scale = np.double(mul_ofm.quantization.scale_f32)
874 alpha_scale, alpha_shift = scaling.elementwise_mul_scale(mul_ifm_scale, mul_ifm2_scale, mul_ofm_scale)
875 op.attrs["alpha_scaling"] = (alpha_scalar, alpha_scale, alpha_shift)
Tim Hall79d07d22020-04-27 18:20:16 +0100876 elif val == -1:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200877 new_op = Op.Abs
Tim Hall79d07d22020-04-27 18:20:16 +0100878 else:
879 return op
880
Louis Verhaardaee5d752020-09-30 09:01:52 +0200881 op.type = new_op
882 op.name = op.name.replace("Maximum", new_op.name)
883 op.outputs[0].name = op.outputs[0].name.replace("Maximum", new_op.name)
Tim Hall79d07d22020-04-27 18:20:16 +0100884 op.inputs = [shared_in]
Tim Halle6ccd872020-11-09 16:46:37 +0000885
886 # Record optimisation in debug database
887 DebugDatabase.add_optimised(op, op)
888
Tim Hall79d07d22020-04-27 18:20:16 +0100889 return op
890
891
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200892def convert_lrelu_to_mul_max(op, arch):
893 # Converts LeakyRelu to Max(alpha * IFM, identity * IFM)
894 # (the opposite of convert_mul_max_to_abs_or_lrelu)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200895 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +0100896 if ifm is None or ofm is None:
897 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200898
899 # Add multiplication with alpha
Louis Verhaardaee5d752020-09-30 09:01:52 +0200900 mul_alpha = Operation(Op.Mul, op.name + "_mul_alpha")
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200901 mul_alpha.add_input_tensor(ifm)
902 # Create const tensor containing alpha as scalar
903 alpha = op.attrs["alpha"]
904 quantization = ifm.quantization.clone()
905 quantization.min = 0
906 quantization.max = alpha * (quantization.quant_max - quantization.quant_min)
907 quantization.scale_f32 = alpha
908 quantization.zero_point = 0
909 alpha_tens = create_const_tensor(op.name + "_alpha_scalar", [], ifm.dtype, [1], np.int8, quantization=quantization)
910 mul_alpha.add_input_tensor(alpha_tens)
911 fm_alpha = ofm.clone(op.name + "_alpha")
912 mul_alpha.set_output_tensor(fm_alpha)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000913 mul_alpha.set_ifm_ofm_shapes()
Tim Halle6ccd872020-11-09 16:46:37 +0000914 DebugDatabase.add_optimised(op, mul_alpha)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200915
Tim Hall93582962020-09-09 21:58:15 +0100916 if check_quantized_tens_scaling_equal(ifm, ofm):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200917 # No identity multiplication is needed
918 fm_id = ifm
919 else:
920 # Add multiplication with identity
Louis Verhaardaee5d752020-09-30 09:01:52 +0200921 mul_identity = Operation(Op.Mul, op.name + "_mul_identity")
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200922 mul_identity.add_input_tensor(ifm)
923 # Create const tensor containing identity as scalar
924 quantization = ifm.quantization.clone()
925 quantization.min = 0
926 quantization.max = quantization.quant_max - quantization.quant_min
927 quantization.scale_f32 = 1
928 quantization.zero_point = 0
929 identity_tens = create_const_tensor(
930 op.name + "_id_scalar", [], ifm.dtype, [1], np.uint8, quantization=quantization
931 )
932 mul_identity.add_input_tensor(identity_tens)
933 fm_id = ofm.clone(op.name + "_id")
934 mul_identity.set_output_tensor(fm_id)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000935 mul_identity.set_ifm_ofm_shapes()
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100936 DebugDatabase.add_optimised(op, mul_identity)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200937
938 # Convert LeakyRelu to Max, add the results of the multiplication(s) as inputs
Louis Verhaardaee5d752020-09-30 09:01:52 +0200939 op.type = Op.Maximum
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200940 op.name = op.name.replace("LeakyRelu", "Maximum")
941 op.inputs = []
942 ifm.consumer_list.remove(op)
943 op.add_input_tensor(fm_alpha)
944 op.add_input_tensor(fm_id)
Tim Halle6ccd872020-11-09 16:46:37 +0000945
946 DebugDatabase.add_optimised(op, op)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200947 return op
948
949
Louis Verhaard2e186c72020-10-09 10:47:04 +0200950def convert_to_lut(op, lut_values, lut_name):
Louis Verhaardf03bad32020-09-25 08:30:44 +0200951 # Rewrite the operation by Add with scalar 0 + LUT activation
952 ifm = op.inputs[0]
Tim Hall93582962020-09-09 21:58:15 +0100953 if ifm is None:
954 return op
Louis Verhaard58520b92020-08-24 16:45:38 +0200955 assert ifm.dtype.size_in_bytes() == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +0200956 op.type = Op.Add
Louis Verhaard2e186c72020-10-09 10:47:04 +0200957 op.name = op.name + "_lut_" + lut_name
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200958 # Mark as no-op to enable potential fusing optimizations
959 op.attrs["is_nop"] = True
960 # Create an input tensor containing scalar zero
961 quantization = QuantizationParameters(0.0, 255.0)
Louis Verhaardd7911c42020-08-25 13:36:41 +0200962 quantization.scale_f32 = ifm.quantization.scale_f32
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200963 quantization.zero_point = 0
Louis Verhaard2e186c72020-10-09 10:47:04 +0200964 tens = create_const_tensor(op.inputs[0].name + "_scalar0", [], ifm.dtype, [0], np.uint8, quantization=quantization)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +0200965 op.add_input_tensor(tens)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000966 op.ifm_shapes.append(Shape4D(tens.shape))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100967
Louis Verhaardf03bad32020-09-25 08:30:44 +0200968 # The LUT must be applied without any preceding rescaling (the LUT itself performs the rescale),
969 # so even if the OFM has a different scale than the IFM, the generated OFM scale instructions
970 # should be the same as the IFM
Louis Verhaardaee5d752020-09-30 09:01:52 +0200971 op.forced_output_quantization = ifm.quantization
Louis Verhaard2e186c72020-10-09 10:47:04 +0200972 lut_tensor = lut.create_lut_tensor(op.name + "_values", lut_values, DataType.int8)
Louis Verhaardf03bad32020-09-25 08:30:44 +0200973 op.set_activation_lut(lut_tensor)
974 return op
975
976
Louis Verhaard2e186c72020-10-09 10:47:04 +0200977def convert_to_lut8(op, fn, fn_name):
Louis Verhaardf03bad32020-09-25 08:30:44 +0200978 # Converts op to a no-op + int8/uint8 LUT which is generated with the given function.
979 # fn is a function(real) -> real
Louis Verhaardaee5d752020-09-30 09:01:52 +0200980 ifm, ofm = op.get_ifm_ofm()
Louis Verhaardf03bad32020-09-25 08:30:44 +0200981 if ifm.dtype not in (DataType.uint8, DataType.int8) or ifm.dtype != ofm.dtype:
982 return op
983 # Generate the LUT
984 ifm_scale = np.double(ifm.quantization.scale_f32)
985 ofm_scale = np.double(ofm.quantization.scale_f32)
986 zp_in = ifm.quantization.zero_point
987 zp_out = ofm.quantization.zero_point
988 values = []
989 ix = range(256) if ifm.dtype == DataType.uint8 else range(-128, 128)
990 quantized_min = min(ix)
991 quantized_max = max(ix)
992 for x in ix:
993 x_real = ifm_scale * (x - zp_in)
994 y_real = fn(x_real)
995 lut_result = round_away_zero(zp_out + y_real / ofm_scale)
996 lut_result = min(quantized_max, max(quantized_min, lut_result))
997 values.append(lut_result)
Louis Verhaard2e186c72020-10-09 10:47:04 +0200998 return convert_to_lut(op, values, fn_name)
Louis Verhaardf03bad32020-09-25 08:30:44 +0200999
1000
1001def convert_lrelu_to_lut(op, arch):
Louis Verhaardaee5d752020-09-30 09:01:52 +02001002 ifm, ofm = op.get_ifm_ofm()
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001003 # Generate the LUT
Louis Verhaardd7911c42020-08-25 13:36:41 +02001004 alpha = op.attrs["alpha"]
1005 ifm_scale = np.double(ifm.quantization.scale_f32)
1006 ofm_scale = np.double(ofm.quantization.scale_f32)
1007 zp_in = ifm.quantization.zero_point
1008 zp_out = ofm.quantization.zero_point
1009 identity_scale, identity_shift = scaling.elementwise_mul_scale(ifm_scale, 1, ofm_scale)
1010 alpha_scalar = 1
1011 alpha_scale, alpha_shift = scaling.elementwise_mul_scale(ifm_scale, alpha, ofm_scale)
1012 if "alpha_scaling" in op.attrs:
1013 # The LeakyRelu was the result from convert_mul_max_to_abs_or_lrelu
1014 alpha_scalar, alpha_scale, alpha_shift = op.attrs["alpha_scaling"]
1015 values = []
Louis Verhaard58520b92020-08-24 16:45:38 +02001016 ix = range(256) if ifm.dtype == DataType.uint8 else range(-128, 128)
Louis Verhaardd7911c42020-08-25 13:36:41 +02001017 quantized_min = min(ix)
1018 quantized_max = max(ix)
1019 for x in ix:
1020 if x < zp_in:
1021 lut_result = zp_out + fp_math.multiply_by_quantized_multiplier(
1022 alpha_scalar * (x - zp_in), alpha_scale, alpha_shift
1023 )
1024 else:
1025 lut_result = zp_out + fp_math.multiply_by_quantized_multiplier(x - zp_in, identity_scale, identity_shift)
1026 lut_result = min(quantized_max, max(quantized_min, lut_result))
1027 values.append(lut_result)
Louis Verhaard2e186c72020-10-09 10:47:04 +02001028 return convert_to_lut(op, values, "lrelu")
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001029
1030
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001031def convert_lrelu(op, arch, nng):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001032 # Converts LeakyRelu to a LUT based solution if possible, otherwise a mul + max
Louis Verhaardaee5d752020-09-30 09:01:52 +02001033 if op.type != Op.LeakyRelu:
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001034 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001035 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +01001036 if ifm is None or ofm is None:
1037 return op
Louis Verhaardd7911c42020-08-25 13:36:41 +02001038 if ifm.dtype in (DataType.uint8, DataType.int8) and ifm.dtype == ofm.dtype:
1039 # use LUT for int8/uint8
1040 return convert_lrelu_to_lut(op, arch)
Tim Hall93582962020-09-09 21:58:15 +01001041 if check_quantized_tens_scaling_equal(ifm, ofm) and ifm.dtype == ofm.dtype == DataType.int16:
Louis Verhaardd7911c42020-08-25 13:36:41 +02001042 # use LeakyRelu unmodified for int16 with equal input/output scaling
1043 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001044 return convert_lrelu_to_mul_max(op, arch)
1045
1046
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001047def convert_tanh_sigmoid_to_lut(op, arch, nng):
Louis Verhaardf03bad32020-09-25 08:30:44 +02001048 # Converts int8/uint8 Sigmoid and Tanh to a LUT based solution
Louis Verhaardaee5d752020-09-30 09:01:52 +02001049 if op.type == Op.Sigmoid:
Louis Verhaard2e186c72020-10-09 10:47:04 +02001050 return convert_to_lut8(op, clamp_sigmoid, "sigmoid")
Louis Verhaardaee5d752020-09-30 09:01:52 +02001051 elif op.type == Op.Tanh:
Louis Verhaard2e186c72020-10-09 10:47:04 +02001052 return convert_to_lut8(op, math.tanh, "tanh")
Louis Verhaardf03bad32020-09-25 08:30:44 +02001053 return op
1054
1055
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001056def remove_unwanted_reshapes(op, arch, nng):
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001057 # Try to remove reshapes enclosing ElementWise operator with only one non-constant input
Louis Verhaardaee5d752020-09-30 09:01:52 +02001058 if not op.run_on_npu or not op.type.is_elementwise_op():
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001059 return op
1060
1061 # Check if the ElementWise operator only have one non-constant input
Louis Verhaardaee5d752020-09-30 09:01:52 +02001062 non_const_tens = [x for x in op.inputs if x.ops[0].type != Op.Const]
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001063 if len(non_const_tens) != 1:
1064 return op
1065 ifm = non_const_tens[0]
1066
1067 # Check if operation is enclosed by Reshapes that can be removed
1068 ofm = op.outputs[0]
1069 prev_op = ifm.ops[0]
1070 if (
1071 len(ifm.consumer_list) == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +02001072 and prev_op.type == Op.Reshape
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001073 and len(ofm.consumer_list) == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +02001074 and ofm.consumer_list[0].type == Op.Reshape
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001075 ):
1076 # Operation is enclosed by reshapes, check if they can be removed
Louis Verhaardaee5d752020-09-30 09:01:52 +02001077 prev_op_ifm, prev_op_ofm = prev_op.get_ifm_ofm()
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001078 cons_op = ofm.consumer_list[0]
1079 cons_op_ifm = ofm
1080 cons_op_ofm = cons_op.outputs[0]
1081 if len(prev_op_ifm.shape) == len(cons_op_ofm.shape):
1082 # Check if quantization is the same in the input and output for the reshape ops
Tim Hall93582962020-09-09 21:58:15 +01001083 if check_quantized_tens_scaling_equal(prev_op_ifm, prev_op_ofm) and check_quantized_tens_scaling_equal(
1084 cons_op_ifm, cons_op_ofm
1085 ):
Patrik Gustavsson7ad862a2020-09-29 14:09:43 +02001086 op.set_input_tensor(prev_op_ifm, 0)
1087 op.set_output_tensor(cons_op_ofm)
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001088 return op
1089
1090
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001091def fuse_activation_function_with_prev(op, arch, nng):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001092 # if op is a no-op: attempts to move the activation function to the preceding op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001093 if not op.attrs.get("is_nop", False) or op.activation is None:
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001094 return op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001095 ifm, ofm = op.get_ifm_ofm()
Tim Hall93582962020-09-09 21:58:15 +01001096 if ifm is None or ofm is None:
1097 return op
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001098 # finds the input(s) to the operation
1099 prev_op = ifm.ops[0]
1100 # Note: the below checks on prev_op require that a first optimize pass on the full graph has been performed
1101 fuse = (
1102 prev_op.run_on_npu
Louis Verhaardaee5d752020-09-30 09:01:52 +02001103 and prev_op.type.npu_block_type != NpuBlockType.Default
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001104 and len(ifm.ops) == 1
1105 and len(prev_op.outputs[0].consumers()) == 1
Louis Verhaardaee5d752020-09-30 09:01:52 +02001106 and prev_op.activation is None
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001107 )
1108 if op.activation_lut is not None and arch.shram_reserved_unused_banks == 0:
1109 # TODO: if SHRAM LUT space is shared with SHRAM ACC (32, 64 MAC),
1110 # LUT currently only works correctly for elementwise ops
1111 fuse = False
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001112 if not fuse:
1113 return op
1114 # Move the fused activation function + corresponding info to prev_op
Louis Verhaardaee5d752020-09-30 09:01:52 +02001115 prev_op.activation = op.activation
1116 prev_op.forced_output_quantization = op.forced_output_quantization
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001117 if op.activation_lut is not None:
1118 prev_op.set_activation_lut(op.activation_lut)
1119 # Bypass op
Louis Verhaard98a34992020-09-01 10:39:04 +02001120 prev_op.set_output_tensor(ofm)
Tim Halle6ccd872020-11-09 16:46:37 +00001121 DebugDatabase.add_optimised(op, prev_op)
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001122 return op
1123
1124
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001125def add_attrs_to_resizebilinear(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +02001126 if op.type == Op.ResizeBilinear and op.run_on_npu:
Dwight Lidman42fed942020-05-29 09:37:03 +02001127 input_tensor = op.inputs[0]
1128 upscaled_shape = [input_tensor.shape[1] * 2, input_tensor.shape[2] * 2]
1129 out_shape = op.outputs[0].shape[1:3]
1130 if not op.attrs["align_corners"] and out_shape == upscaled_shape:
1131 # this means the output is supposed to be a x2 upscale,
1132 # so we need to do SAME padding
Michael McGeagh16895482020-12-14 15:51:20 +00001133 op.attrs["padding"] = Padding.SAME
Dwight Lidman42fed942020-05-29 09:37:03 +02001134 elif op.attrs["align_corners"] and out_shape == [upscaled_shape[0] - 1, upscaled_shape[1] - 1]:
1135 # here we can just run the avg pool without padding and
1136 # produce a (M * 2 - 1, N * 2 - 1) sized output
Michael McGeagh16895482020-12-14 15:51:20 +00001137 op.attrs["padding"] = Padding.VALID
Dwight Lidman42fed942020-05-29 09:37:03 +02001138 else:
Charles Xu9a03fdf2020-07-02 15:12:40 +02001139 return op
Dwight Lidman42fed942020-05-29 09:37:03 +02001140 input_tensor.resampling_mode = resampling_mode.NEAREST
Tim Hallc30f4952020-06-15 20:47:35 +01001141 op.attrs.update({"strides": (1, 1, 1, 1), "ksize": (1, 2, 2, 1)})
Dwight Lidman42fed942020-05-29 09:37:03 +02001142 return op
1143
1144
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001145def fixup_bias_tensors(op, arch, nng):
Louis Verhaardaee5d752020-09-30 09:01:52 +02001146 if op.type.needs_bias() and op.bias is None:
Jacob Bohlina41cd4d2020-08-26 18:21:28 +02001147 # Op has no bias, add bias tensor filled with zeros
1148 nr_biases = op.inputs[1].shape[-1]
1149 bias_values = [0] * nr_biases
1150 bias_tensor = create_const_tensor(op.name + "_bias", [nr_biases], DataType.int32, bias_values)
1151 bias_tensor.quant_values = bias_tensor.values
1152 op.set_input_tensor(bias_tensor, -1)
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +02001153
1154 return op
1155
1156
Patrik Gustavsson3010d9b2020-10-01 08:22:10 +02001157def supported_operator_check(op, arch, nng):
Tim Hall79d07d22020-04-27 18:20:16 +01001158 op.run_on_npu = arch.supported_operators.is_operator_supported(op)
1159 return op
1160
1161
Tim Halle6ccd872020-11-09 16:46:37 +00001162def _record_optimised(op, arch):
1163 if op.type != Op.Const:
1164 DebugDatabase.add_optimised(op, op)
1165
1166
Tim Hall79d07d22020-04-27 18:20:16 +01001167def optimise_graph_a(nng, arch, verbose_graph=False):
1168 if verbose_graph:
1169 nng.print_graph()
1170
Patrik Gustavsson2349d422020-12-01 16:02:29 +01001171 pre_process_list = [
1172 supported_operator_check,
1173 set_ifm_ofm_op_shapes,
1174 # TODO: memory-only Op removal
1175 ]
1176
1177 for idx, sg in enumerate(nng.subgraphs):
1178 # rewrite graph pass
1179 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
1180 nng, sg, arch, [], pre_process_list, rewrite_unsupported=False,
1181 )
1182
Tim Hall79d07d22020-04-27 18:20:16 +01001183 op_rewrite_list = [
Tim Hall4e127762020-05-15 16:05:49 +01001184 set_tensor_equivalence,
Tim Hall79d07d22020-04-27 18:20:16 +01001185 convert_depthwise_to_conv,
Michael McGeagh8d939c02020-07-29 13:11:43 +01001186 convert_conv_to_fc,
Fredrik Svedberga0c36242020-06-03 15:43:31 +02001187 convert_softmax,
Diqing Zhong016b8272020-12-16 16:46:06 +01001188 optimise_strided_conv,
Tim Hall79d07d22020-04-27 18:20:16 +01001189 fixup_fully_connected_input,
Diqing Zhong94457b12020-12-09 15:22:40 +01001190 convert_batched_fc_shape,
Tim Hall79d07d22020-04-27 18:20:16 +01001191 fixup_pack_input,
Fredrik Svedberg0f98b362020-09-29 10:00:39 +02001192 unfuse_activation_function,
Tim Hall79d07d22020-04-27 18:20:16 +01001193 fixup_conv2d_backprop,
Michael McGeagh8dbf8cf2020-09-08 11:09:48 +01001194 fixup_relus_with_differing_ifm_ofm_scaling,
Tim Hall79d07d22020-04-27 18:20:16 +01001195 fixup_act_reorder,
Charles Xu78792222020-05-13 10:15:26 +02001196 fixup_elementwise_with_scalars,
Jacob Bohline843d332020-06-23 12:12:56 +02001197 reorder_depthwise_weights,
Charles Xu9a03fdf2020-07-02 15:12:40 +02001198 fixup_resizebilinear,
Jacob Bohlina41cd4d2020-08-26 18:21:28 +02001199 fixup_bias_tensors,
Dwight Lidmanc3862c22020-09-14 15:22:33 +02001200 convert_nop_split_to_identity,
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001201 convert_mul_max_to_abs_or_lrelu,
Patrik Gustavssonfa4cb292020-09-10 08:19:36 +02001202 remove_unwanted_reshapes,
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001203 convert_lrelu,
Louis Verhaardf03bad32020-09-25 08:30:44 +02001204 convert_tanh_sigmoid_to_lut,
Tim Hall79d07d22020-04-27 18:20:16 +01001205 ]
1206
1207 for idx, sg in enumerate(nng.subgraphs):
1208 # rewrite graph pass
1209 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
Dwight Lidman73320a42020-11-05 10:34:41 +01001210 nng, sg, arch, [], op_rewrite_list, rewrite_unsupported=False,
Tim Hall79d07d22020-04-27 18:20:16 +01001211 )
1212
1213 for idx, sg in enumerate(nng.subgraphs):
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001214 # remove passthrough tensors and attempt further optimizations
1215 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
Patrik Gustavsson2349d422020-12-01 16:02:29 +01001216 nng, sg, arch, [remove_passthrough_tensor], [fuse_activation_function_with_prev, add_padding_fields],
Louis Verhaardb9fc33c2020-08-13 11:47:36 +02001217 )
Tim Hall79d07d22020-04-27 18:20:16 +01001218
Tim Halle6ccd872020-11-09 16:46:37 +00001219 # Post-optimisation operator debug tracing
1220 for sg in nng.subgraphs:
1221 rewrite_graph.visit_graph_post_order(sg.output_tensors, arch, [], [_record_optimised])
1222
Tim Hall79d07d22020-04-27 18:20:16 +01001223 if verbose_graph:
1224 nng.print_graph()
1225 return nng
1226
Diego Russoea6111a2020-04-14 18:41:58 +01001227
Tim Hall79d07d22020-04-27 18:20:16 +01001228def optimise_graph_b(nng, arch, verbose_graph=False):
1229 if verbose_graph:
1230 nng.print_graph()
1231
1232 for idx, sg in enumerate(nng.subgraphs):
1233 # combined rewrite graph pass
Dwight Lidmanc6ac1942020-10-02 14:55:45 +02001234 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
patrik.gustavssoneeb85152020-12-21 17:10:40 +00001235 nng, sg, arch, [fixup_unpack_output, fixup_stridedslice_output, rewrite_concat, rewrite_split], [],
Dwight Lidmanc6ac1942020-10-02 14:55:45 +02001236 )
Tim Hall79d07d22020-04-27 18:20:16 +01001237
1238 if verbose_graph:
1239 nng.print_graph()
1240 return nng