blob: 3fe703e190df3bd20f764cd634153a01217c2e3c [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
Diego Russoea6111a2020-04-14 18:41:58 +010020
21import numpy as np
22
23from . import rewrite_graph
Diego Russoea6111a2020-04-14 18:41:58 +010024from .data_type import DataType
Louis Verhaard7db78962020-05-25 15:05:26 +020025from .errors import UnsupportedFeatureError
Dwight Lidman42fed942020-05-29 09:37:03 +020026from .ethos_u55_regs.ethos_u55_regs import resampling_mode
Louis Verhaarde0ef2732020-06-03 08:56:44 +020027from .numeric_util import full_shape
Diego Russoe8a10452020-04-21 17:39:10 +010028from .operation import NpuBlockType
29from .operation import Operation
Fredrik Svedberga0c36242020-06-03 15:43:31 +020030from .softmax import SoftMax
Michael McGeaghc5b549b2020-08-07 11:54:28 +010031from .tensor import create_const_tensor
32from .tensor import create_reshape_tensor
Charles Xu9a03fdf2020-07-02 15:12:40 +020033from .tensor import QuantizationParameters
Diego Russoe8a10452020-04-21 17:39:10 +010034from .tensor import Tensor
Tim Hall79d07d22020-04-27 18:20:16 +010035
36passthrough_nodes = set(("Identity",))
37
38
39def remove_passthrough_tensor(tens, arch):
40 if len(tens.ops) == 1 and tens.ops[0].type in passthrough_nodes:
41 assert len(tens.ops[0].inputs) == 1
42 tens = tens.ops[0].inputs[0]
43 return tens
44
45
46def rewrite_concat(tens, arch):
47 if len(tens.ops) == 1 and tens.ops[0].is_concat_op():
48 concat_op = tens.ops[0]
49 if tens != concat_op.outputs[0]:
50 return tens # don't attempt to rewrite the min/max outputs of QuantizedConcat
51
52 # Not supported so leave it and run on CPU
53 if not concat_op.run_on_npu:
54 return tens
55
56 inputs, axis = concat_op.get_concat_inputs_axis()
57
58 tens.ops = []
59 offset = 0
60 for idx, inp in enumerate(inputs):
61 new_op = Operation("ConcatSliceWrite", concat_op.name + str(idx))
62 new_op.inputs = [inp]
63 new_op.outputs = [tens]
64 new_op.attrs["concat_axis"] = axis
65 new_op.attrs["concat_start"] = offset
66 offset += inp.shape[axis]
67 new_op.attrs["concat_end"] = offset
68 new_op.run_on_npu = True
69 tens.ops.append(new_op)
70 assert tens.shape[axis] == offset
71
Patrik Gustavsson458a2082020-08-13 13:41:05 +020072 # If axis = 3, NHCWB16 can only be used in the output if all the concat_start's are a multiple of 16,
73 # as it is only then the address offset for the ofm, for all operations, will be 16 byte aligned
74 # For other values of axis the address offsets will be 16 byte aligned, as they are all based on c = 0
75 # and those addresses are always 16 byte aligned due to the NHCWB16 format.
76 if axis == 3:
77 for op in tens.ops:
78 if op.attrs["concat_start"] % 16 != 0:
79 tens.avoid_NHCWB16 = True
80 break
81
Tim Hall79d07d22020-04-27 18:20:16 +010082 return tens
83
84
85def rewrite_split(tens, arch):
86
87 if len(tens.ops) == 1 and tens.ops[0].is_split_op():
88 split_op = tens.ops[0]
89
90 # Not supported so leave it and run on CPU
91 if not split_op.run_on_npu:
92 return tens
93
94 inp, outputs, axis, offset_start, offset_end = split_op.get_split_inputs_axis()
95
96 tens.ops = []
97 new_op = Operation("SplitSliceRead", split_op.name)
98 new_op.inputs = [inp]
Tim Hall79d07d22020-04-27 18:20:16 +010099
100 # For Split the offset cannot be extracted from the tensor so it has to
101 # be calculated from the index of the output tensor
Diego Russoea6111a2020-04-14 18:41:58 +0100102 if axis is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100103 # Get the start and end of the split
104 offset_start = [0] * len(tens.shape)
105 offset_end = [0] * len(tens.shape)
106 for out in outputs:
107 if out == tens:
108 break
109 offset_start[axis] += out.shape[axis]
110
111 offset_end[axis] = offset_start[axis] + tens.shape[axis]
112
113 new_op.attrs["split_start"] = offset_start
114 new_op.attrs["split_end"] = offset_end
115 new_op.run_on_npu = True
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100116 new_op.set_output_tensor(tens)
Tim Hall79d07d22020-04-27 18:20:16 +0100117
118 return tens
119
120
121def needed_total_padding(input_size, stride, filter_size):
122 out_size = (input_size + stride - 1) // stride
123 needed_input = (out_size - 1) * stride + filter_size
124 total_padding = max(0, needed_input - input_size)
125 return total_padding
126
127
128def calc_padding_and_skirt(padding_type, kernel_size, stride, input_dims):
129 ypad = needed_total_padding(int(input_dims[1]), int(stride[1]), int(kernel_size[0]))
130 xpad = needed_total_padding(int(input_dims[2]), int(stride[2]), int(kernel_size[1]))
131 if padding_type == b"SAME":
132 left_pad = (xpad + 0) // 2
133 right_pad = (xpad + 1) // 2
134 top_pad = (ypad + 0) // 2
135 bottom_pad = (ypad + 1) // 2
136 elif padding_type == b"VALID":
137 left_pad = 0
138 right_pad = 0
139 top_pad = 0
140 bottom_pad = 0
141 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200142 raise UnsupportedFeatureError("Unknown padding {}".format(str(padding_type)))
Tim Hall79d07d22020-04-27 18:20:16 +0100143 padding = (top_pad, left_pad, bottom_pad, right_pad)
144 skirt = (top_pad, left_pad, ypad - top_pad, xpad - left_pad)
145 return padding, skirt
146
Tim Hallc30f4952020-06-15 20:47:35 +0100147
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200148def calc_upscaled_padding_and_skirt(padding_type, kernel_size, stride, input_dims, upscaling_factor):
149 kernel_height, kernel_width = kernel_size[0], kernel_size[1]
Jacob Bohlincf7da102020-05-20 09:03:40 +0200150 if padding_type == b"SAME":
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200151 ypad = needed_total_padding(int(input_dims[1]) * upscaling_factor, int(stride[1]), int(kernel_height))
152 xpad = needed_total_padding(int(input_dims[2]) * upscaling_factor, int(stride[2]), int(kernel_width))
153
154 right_pad = ((xpad + 1) // upscaling_factor) - 1
155 bottom_pad = ((ypad + 1) // upscaling_factor) - 1
156 left_pad = max(kernel_width - 1 - right_pad, 0)
157 top_pad = max(kernel_height - 1 - bottom_pad, 0)
158
Jacob Bohlincf7da102020-05-20 09:03:40 +0200159 elif padding_type == b"VALID":
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200160 right_pad = max(kernel_width - 2, 0)
161 bottom_pad = max(kernel_height - 2, 0)
162 left_pad = kernel_width - 1
163 top_pad = kernel_height - 1
Jacob Bohlincf7da102020-05-20 09:03:40 +0200164 else:
165 assert 0, "Unknown padding"
166
167 padding = (top_pad, left_pad, bottom_pad, right_pad)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200168 skirt = padding
Jacob Bohlincf7da102020-05-20 09:03:40 +0200169 return padding, skirt
170
Tim Hall79d07d22020-04-27 18:20:16 +0100171
172def fixup_conv2d_backprop(op, arch):
173 if op.type == "Conv2DBackpropInput":
174 # flip the inputs
175 op.inputs[0], op.inputs[2] = op.inputs[2], op.inputs[0]
Jacob Bohlincf7da102020-05-20 09:03:40 +0200176 op.type = "Conv2DBackpropInputSwitchedBias"
177 weight_shape = op.inputs[1].shape
178 weight_sets = weight_shape[3]
179
180 if len(op.inputs) < 4:
181 # Add bias/scale tensor filled with zeros
Jacob Bohlincf7da102020-05-20 09:03:40 +0200182 scale_tens = Tensor([weight_sets], DataType.int32, op.name + "_bias_tens")
183 scale_tens.values = [0] * weight_sets
184 scale_tens.quant_values = [0] * weight_sets
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100185 scale_op = Operation("Const", op.name + "_bias")
186 scale_op.set_output_tensor(scale_tens)
187 op.add_input_tensor(scale_tens)
Jacob Bohlincf7da102020-05-20 09:03:40 +0200188
189 # Update strides
Tim Hallc30f4952020-06-15 20:47:35 +0100190 op.attrs.update({"stride_w": 1, "stride_h": 1, "strides": (1, 1, 1, 1)})
Tim Hall79d07d22020-04-27 18:20:16 +0100191
192 return op
193
194
Charles Xu9a03fdf2020-07-02 15:12:40 +0200195# Convert the op to an elementwise add
196def convert_resizebilinear_1x1_to_add(op):
197 op.type = "AddAct"
198 op.name = op.name + "_add"
199 op.attrs.update({"npu_block_type": NpuBlockType.ElementWise})
200 op.attrs["resizebilinear"] = True
201 # Create an input tensor filled with zeros
202 shape = op.outputs[0].shape
203 tens = Tensor(shape, op.inputs[0].dtype, op.inputs[1].name + "_add")
204 tens.values = np.zeros(shape)
205 tens.quant_values = np.zeros(shape, np.uint8)
206 tens.quantization = QuantizationParameters(0.0, 255.0)
207 tens.quantization.scale_f32 = 1.0
208 tens.quantization.zero_point = 0
209 tens.consumer_list = [op]
210 tens_op = op.inputs[1].ops[0]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100211 tens_op.set_output_tensor(tens)
Charles Xu9a03fdf2020-07-02 15:12:40 +0200212 # Set the add inputs
213 op.inputs[1] = op.inputs[0]
214 op.inputs[0] = tens
215
216 return op
217
218
219def fixup_resizebilinear(op, arch):
220 if op.type == "ResizeBilinear":
221 if op.inputs[0].shape[1] == 1 and op.inputs[0].shape[2] == 1:
222 convert_resizebilinear_1x1_to_add(op)
Charles Xu36ffaf32020-08-05 15:40:44 +0200223 elif op.inputs[0].shape == op.outputs[0].shape:
224 # Bypass nop resizebilinear
225 op.inputs = op.inputs[:1]
226 op.type = "Identity"
Charles Xu9a03fdf2020-07-02 15:12:40 +0200227
228 return op
229
230
Tim Hall79d07d22020-04-27 18:20:16 +0100231def fixup_fully_connected_input(op, arch):
232 if op.type == "FullyConnectedAct":
233 inp = op.inputs[0]
234 weights = op.inputs[1]
235
236 n_in_elems = weights.shape[-2]
237 elms = inp.elements()
238 batch_size = elms // n_in_elems
239 assert batch_size * n_in_elems == elms
240
241 desired_shape = [batch_size, n_in_elems]
242 if inp.shape != desired_shape:
243 # mismatch, insert a reshape to fix this.
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100244 op.inputs[0] = create_reshape_tensor(inp, desired_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100245
246 return op
247
248
249def fixup_pack_input(op, arch):
250 if op.type == "Pack":
251 # Pack is also referred to as Stack
252 # Requires the rewrite_concat function to be called on the op afterwards
253 axis = int(op.attrs["axis"])
254 desired_shape = op.inputs[0].shape[:axis] + [1] + op.inputs[0].shape[axis:]
255
256 # Construct 1 shape tensor to be used by all inserted reshape ops
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100257 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, desired_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100258
259 for idx, inp in enumerate(op.inputs):
Tim Hall79d07d22020-04-27 18:20:16 +0100260 reshape_out = inp.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100261 reshape_out.set_all_shapes(desired_shape)
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100262
263 reshape_op = Operation("Reshape", "{}{}_reshape".format(op.name, idx))
264 reshape_op.attrs["new_shape"] = desired_shape
265 reshape_op.inputs = [inp, new_shape_tens]
266 reshape_op.set_output_tensor(reshape_out)
Tim Hall79d07d22020-04-27 18:20:16 +0100267
268 op.inputs[idx] = reshape_out
269
270 op.type = "PackReshaped"
271
272 return op
273
274
275def fixup_unpack_output(tens, arch):
276 op = tens.ops[0]
277 if op.type in set(("Unpack", "StridedSlice")):
278 # Unpack is also referred to as Unstack
279 # Requires the rewrite_split function to be called on the op afterwards
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200280
281 reshape_input_shape = tens.shape
Tim Hall79d07d22020-04-27 18:20:16 +0100282 if op.type == "StridedSlice":
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200283 new_axis_mask = op.attrs["new_axis_mask"]
Tim Hall79d07d22020-04-27 18:20:16 +0100284 shrink_axis_mask = op.attrs["shrink_axis_mask"]
Louis Verhaard7db78962020-05-25 15:05:26 +0200285 ellipsis_mask = op.attrs["ellipsis_mask"]
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200286
287 if (new_axis_mask != 0 and shrink_axis_mask != 0) or ellipsis_mask != 0:
288 # Not supported, will be put on CPU
289 return tens
290 if shrink_axis_mask == 0 and new_axis_mask == 0:
Tim Hall79d07d22020-04-27 18:20:16 +0100291 # Equal Rank StridedSlice, no need to insert reshape
292 return tens
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200293 elif shrink_axis_mask != 0:
294 n = 0
295 axis = 0
296 while shrink_axis_mask:
297 prev_mask = shrink_axis_mask
298 n += 1
299 shrink_axis_mask &= shrink_axis_mask - 1
300 axis = int(math.log2(prev_mask - shrink_axis_mask))
301 reshape_input_shape = reshape_input_shape[:axis] + [1] + reshape_input_shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100302
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200303 assert len(tens.shape) == (len(op.inputs[0].shape) - n)
304 op.attrs["shrink_axis_mask"] = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100305
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200306 elif new_axis_mask != 0:
307 n = 0
308 axis = 0
309 while new_axis_mask:
310 prev_mask = new_axis_mask
311 n += 1
312 new_axis_mask &= new_axis_mask - 1
313 axis = int(math.log2(prev_mask - new_axis_mask))
Louis Verhaard7db78962020-05-25 15:05:26 +0200314 reshape_input_shape = reshape_input_shape[:axis] + reshape_input_shape[(axis + 1) :]
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200315 new_axis_mask >>= 1
316
317 assert len(tens.shape) == (len(op.inputs[0].shape) + n)
318 op.attrs["new_axis_mask"] = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100319 else:
320 axis = int(op.attrs["axis"])
321 op.type = "UnpackReshaped"
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200322 reshape_input_shape = tens.shape[:axis] + [1] + tens.shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100323
324 # Construct 1 shape tensor to be used by all inserted reshape ops
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100325 new_shape_tens = create_const_tensor(op.name + "_reshape_shape", [1], DataType.int32, tens.shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100326
327 for idx, out_tens in enumerate(op.outputs):
Tim Hall79d07d22020-04-27 18:20:16 +0100328 reshape_in = out_tens.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100329 reshape_in.set_all_shapes(reshape_input_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100330 reshape_in.ops = [op]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100331
332 reshape_op = Operation("Reshape", "{}{}_reshape".format(op.name, idx))
333 reshape_op.attrs["new_shape"] = reshape_input_shape
Tim Hall79d07d22020-04-27 18:20:16 +0100334 reshape_op.inputs = [reshape_in, new_shape_tens]
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100335 reshape_op.set_output_tensor(out_tens)
Tim Hall79d07d22020-04-27 18:20:16 +0100336
337 op.outputs[idx] = reshape_in
338
339 return tens
340
341
342def add_padding_fields(op, arch):
343 if "padding" in op.attrs:
344 if "Conv" in op.type:
345 kernel_size = op.inputs[1].shape[:2]
346 input_shape = op.inputs[0].shape
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200347 elif "Pool" in op.type or op.type in ("ResizeBilinear", "ReduceSum"):
Tim Hall79d07d22020-04-27 18:20:16 +0100348 kernel_size = op.attrs["ksize"][1:3]
349 input_shape = op.inputs[0].shape
350 elif op.type == "ExtractImagePatches":
351 kernel_size = op.attrs["ksizes"][1:3]
352 input_shape = op.inputs[0].shape
353 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200354 raise UnsupportedFeatureError("Unknown operation that uses padding: {}".format(op.type))
Tim Hall79d07d22020-04-27 18:20:16 +0100355
Jacob Bohlincf7da102020-05-20 09:03:40 +0200356 if op.type == "Conv2DBackpropInputSwitchedBias":
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200357 upscaling_factor = op.outputs[0].shape[1] // input_shape[1]
Tim Hallc30f4952020-06-15 20:47:35 +0100358 padding, skirt = calc_upscaled_padding_and_skirt(
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200359 op.attrs["padding"], kernel_size, op.attrs["strides"], input_shape, upscaling_factor
Tim Hallc30f4952020-06-15 20:47:35 +0100360 )
Jacob Bohlincf7da102020-05-20 09:03:40 +0200361 else:
362 dilation_h, dilation_w = op.get_dilation_h_w()
363 dilated_kernel_size = [dilation_h * (kernel_size[0] - 1) + 1, dilation_w * (kernel_size[1] - 1) + 1]
Tim Hallc30f4952020-06-15 20:47:35 +0100364 padding, skirt = calc_padding_and_skirt(
365 op.attrs["padding"], dilated_kernel_size, op.attrs["strides"], input_shape
366 )
Jacob Bohlincf7da102020-05-20 09:03:40 +0200367
Tim Hall79d07d22020-04-27 18:20:16 +0100368 op.attrs["explicit_padding"] = padding
369 op.attrs["skirt"] = skirt
Jacob Bohlincf7da102020-05-20 09:03:40 +0200370
Tim Hall79d07d22020-04-27 18:20:16 +0100371 return op
372
373
Jacob Bohlincf7da102020-05-20 09:03:40 +0200374conv_op = set(("Conv2D", "QuantizedConv2D", "Conv2DBackpropInputSwitchedBias", "Conv2DBiasAct"))
Tim Hall79d07d22020-04-27 18:20:16 +0100375fc_op = set(
376 (
377 "MatMul",
378 "QuantizedMatMul",
379 "BlockLSTM",
380 "RnnAct",
381 "UnidirectionalSequenceRnnAct",
382 "BidirectionalSequenceRnnAct",
383 "LstmAct",
384 "UnidirectionalSequenceLstmAct",
385 "BidirectionalSequenceLstmAct",
386 "FullyConnectedAct",
387 )
388)
389depthwise_op = set(("DepthwiseConv2dNative", "DepthwiseConv2dBiasAct",))
Louis Verhaard7db78962020-05-25 15:05:26 +0200390pool_op = set(
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200391 ("AvgPool", "MaxPool", "QuantizedAvgPool", "QuantizedMaxPool", "AvgPoolAct", "MaxPoolAct", "ResizeBilinear")
Louis Verhaard7db78962020-05-25 15:05:26 +0200392)
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200393reduce_sum_ops = set(("ReduceSum",))
394elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum", "LeakyRelu", "Abs", "CLZ", "SHL", "SHR"))
Charles Xu78792222020-05-13 10:15:26 +0200395binary_elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum"))
Tim Hall79d07d22020-04-27 18:20:16 +0100396activation_ops = set(("Relu", "Relu6", "ReluN1To1", "Sigmoid", "Tanh"))
397memory_only_ops = set(("Reshape",))
398
Diego Russoea6111a2020-04-14 18:41:58 +0100399
Tim Hall79d07d22020-04-27 18:20:16 +0100400# Check if the op can be reordered
401def get_prepend_op(op):
402 inp = op.inputs[0]
403 # The op should be reordered between prev_op and prep_op
404 prev_op = inp.ops[-1]
405 prep_op = None
406 while prev_op.type in memory_only_ops and len(prev_op.outputs) == 1 and len(prev_op.outputs[0].consumers()) == 1:
407 prep_op = prev_op
408 inp = prev_op.inputs[0]
409 prev_op = inp.ops[-1]
Diego Russoea6111a2020-04-14 18:41:58 +0100410 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 +0100411 return prep_op
412
413 return None
414
415
416def mark_npu_block_type(op, arch):
417 npu_block_type = NpuBlockType.Default
418 if op.type in conv_op:
419 npu_block_type = NpuBlockType.ConvolutionMxN
420 elif op.type in fc_op:
421 npu_block_type = NpuBlockType.VectorProduct
422 elif op.type in depthwise_op:
423 npu_block_type = NpuBlockType.ConvolutionDepthWise
424 elif op.type in pool_op:
425 npu_block_type = NpuBlockType.Pooling
426 elif op.type in elementwise_op:
427 npu_block_type = NpuBlockType.ElementWise
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200428 elif op.type in reduce_sum_ops:
429 npu_block_type = NpuBlockType.ReduceSum
Tim Hall79d07d22020-04-27 18:20:16 +0100430
431 op.attrs["npu_block_type"] = npu_block_type
432 return op
433
434
435def convert_depthwise_to_conv(op, arch):
436 # Depthwise is equivalent to a single conv2d if the ifm depth is 1 and
437 # the ofm depth equals the depth multipler.
438 # If those conditions are true, then we can perform a simple
439 # switch of the operator type (and weight order)
440
441 if ("DepthwiseConv2d" in op.type) and (op.attrs["depth_multiplier"] != 1):
442 ifm_tensor = op.inputs[0]
443 weight_tensor = op.inputs[1]
444 ofm_tensor = op.outputs[0]
445 if (ifm_tensor.shape[3] == 1) and (ofm_tensor.shape[3] == op.attrs["depth_multiplier"]):
446 # Change op type to Conv2d
447 op.type = op.type.replace("DepthwiseConv2d", "Conv2D")
448 del op.attrs["channel_multiplier"]
449 del op.attrs["depth_multiplier"]
450
451 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100452 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Tim Hall79d07d22020-04-27 18:20:16 +0100453 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200454 raise UnsupportedFeatureError(
455 "Unsupported DepthwiseConv2d with depth_multiplier = {}, ifm channels = {}, ofm channels = {}".format(
Tim Hall79d07d22020-04-27 18:20:16 +0100456 op.attrs["depth_multiplier"], ifm_tensor.shape[3], ofm_tensor.shape[3]
457 )
458 )
Tim Hall79d07d22020-04-27 18:20:16 +0100459 return op
460
461
Jacob Bohline843d332020-06-23 12:12:56 +0200462def reorder_depthwise_weights(op, arch):
463 if "DepthwiseConv2d" in op.type:
464 weight_tensor = op.inputs[1]
465 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100466 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Jacob Bohline843d332020-06-23 12:12:56 +0200467 weight_tensor.weight_transpose_depthwise = True
468
469 return op
470
471
Michael McGeagh8d939c02020-07-29 13:11:43 +0100472def convert_conv_to_fc(op, arch):
473 # Conv 1x1 can be equivalent to Fully Connected.
474 # By representing certain convs as fully connected layers, Vela can better determine wether or not to use
475 # caching/double buffering for the weights.
476 # (Weights dont need to be reloaded for convs when IFM H and W are 1)
477 if op.type == "Conv2DBiasAct":
478 _, h, w, _ = op.inputs[0].shape
479 kh, kw, _, _ = op.inputs[1].shape
480 if h == 1 and w == 1 and kh == 1 and kw == 1:
481 # Overwrite this op as a Fully Connected Op
482 op.name += "_fc"
483 op.type = "FullyConnectedAct"
484 faf = op.attrs.get("fused_activation_function", None)
485 op.attrs = {
486 "fused_activation_function": faf,
487 "weights_format": 0,
488 "npu_block_type": NpuBlockType.VectorProduct,
489 }
490 # Reshape Weights to be 2D. HWIO becomes just IO (as H and W are 1, they can just be dropped)
491 weight_tensor = op.inputs[1]
492 weight_tensor.quant_values = weight_tensor.quant_values.squeeze(axis=(0, 1))
493 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
494 # The output from a fully connected is expected to be 2D so we need to add a reshape layer to convert it
495 # back to 4D afterwards as the next layer is expecting that shape
496 orig_ofm_tensor = op.outputs[0]
497 # 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})
498 fc_ofm_tensor = orig_ofm_tensor.clone("_fc")
499 fc_ofm_tensor.set_all_shapes([1, fc_ofm_tensor.shape[-1]])
500 fc_ofm_tensor.ops = [op]
501 # Add a reshape after the new OFM to convert it back to the original 4D shape
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100502 reshape_name = op.name + "_reshape"
503 new_shape_tens = create_const_tensor(reshape_name + "_shape", [1], DataType.int32, orig_ofm_tensor.shape)
Michael McGeagh8d939c02020-07-29 13:11:43 +0100504 reshape_op = Operation("Reshape", reshape_name)
Michael McGeagh8d939c02020-07-29 13:11:43 +0100505 reshape_op.attrs["new_shape"] = orig_ofm_tensor.shape
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100506 reshape_op.inputs = [fc_ofm_tensor, new_shape_tens]
507 reshape_op.set_output_tensor(orig_ofm_tensor)
Michael McGeagh8d939c02020-07-29 13:11:43 +0100508 # Replace this ops OFM to point to the 2D tensor
509 op.outputs[0] = fc_ofm_tensor
510 return op
511
512
Tim Hall79d07d22020-04-27 18:20:16 +0100513# Reorder activation op if it's after the memory only operations
514def fixup_act_reorder(op, arch):
515 if op.type in activation_ops:
516 prep_op = get_prepend_op(op)
Diego Russoea6111a2020-04-14 18:41:58 +0100517 if prep_op is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100518 act_op = op.clone("_reordered")
519 act_op.inputs = [prep_op.inputs[0]]
520 act_op_out = act_op.inputs[0].clone("_acted")
521 act_op_out.quantization = op.outputs[0].quantization.clone()
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100522 act_op.set_output_tensor(act_op_out)
Tim Hall79d07d22020-04-27 18:20:16 +0100523 prep_op.inputs[0] = act_op_out
524 prep_op.outputs[0].quantization = act_op_out.quantization.clone()
525
526 # Mark the op so that it will be removed as passthrough later on
527 op.type = "Identity"
528 return op
529
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200530
Charles Xu78792222020-05-13 10:15:26 +0200531def fixup_elementwise_with_scalars(op, arch):
532 if op.type in binary_elementwise_op:
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200533 ifm_tensor, ifm2_tensor, _, _ = op.get_ifm_ifm2_weights_ofm()
Charles Xu78792222020-05-13 10:15:26 +0200534 if ifm2_tensor.shape != [] and ifm_tensor.shape != []:
535 diff = len(ifm_tensor.shape) - len(ifm2_tensor.shape)
536 if diff > 0:
537 ifm2_tensor.shape = full_shape(len(ifm_tensor.shape), ifm2_tensor.shape, 1)
538 elif diff < 0:
539 ifm_tensor.shape = full_shape(len(ifm2_tensor.shape), ifm_tensor.shape, 1)
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200540 elif ifm_tensor.shape == [] and ifm_tensor.quant_values is None:
541 # IFM is marked as a scalar, but is a result of an operation; change it to a shape of size 1
542 ifm_tensor.shape = len(ifm2_tensor.shape) * [1]
543 ifm_tensor.storage_shape = ifm_tensor.shape
544 elif ifm2_tensor.shape == [] and ifm2_tensor.quant_values is None:
545 # IFM2 is marked as a scalar, but is a result of an operation; change it to a shape of size 1
546 ifm2_tensor.shape = len(ifm_tensor.shape) * [1]
547 ifm2_tensor.storage_shape = ifm2_tensor.shape
Charles Xu78792222020-05-13 10:15:26 +0200548 return op
Tim Hall79d07d22020-04-27 18:20:16 +0100549
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200550
Tim Hall4e127762020-05-15 16:05:49 +0100551# Set input/output tensor equivalence to the same id for memory operations
552def set_tensor_equivalence(op, arch):
553 if op.type == "Reshape":
554 eid = op.outputs[0].equivalence_id
555 for inp in op.inputs:
556 inp.equivalence_id = eid
557 return op
558
559
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200560def convert_softmax(op, arch):
561 if op.type == "Softmax" and op.run_on_npu:
562 softmax = SoftMax(op)
563 op = softmax.get_graph()
564 return op
565
566
Tim Hall79d07d22020-04-27 18:20:16 +0100567def convert_mul_max_to_abs_or_lrelu(op, arch):
Diego Russoea6111a2020-04-14 18:41:58 +0100568 r"""Whenever there is a subgraph with this topology:
Tim Hall79d07d22020-04-27 18:20:16 +0100569
570 Input X For X = -1 or X > 0
571 | \ / This subgraph can be replaced with either
572 | Mul an Abs (if X = -1) or a LeakyReLU (if X > 0)
573 | /
574 Max
575 """
576
577 if op.type == "Maximum":
578 # finds the Mul input(s) to the Max
579 muls = [i for i in op.inputs if i.ops[0].type == "MulAct"]
580 if len(muls) == 1:
581 mul = muls[0].ops[0]
582 elif len(muls) == 2:
583 # In the case both inputs are Muls, find the one with the same input as the Max
584 mul = [m for m in muls if len(set(op.inputs + m.ops[0].inputs)) == 1][0].ops[0]
585 else:
586 # No Mul inputs
587 return op
588
589 # make sure the Mul doesn't have any other consumers
590 if len(mul.outputs[0].consumers()) != 1:
591 return op
592 # make sure the Mul doesn't have a faf
593 if mul.attrs["fused_activation_function"]:
594 return op
595
596 # finds the branched input that goes to both the Max and the Mul
597 shared = set(op.inputs) & set(mul.inputs)
598 if len(shared) == 1:
599 shared_in = shared.pop()
600 # find the constant scalar input to the Mul
601 const_tens = (set(mul.inputs) - {shared_in}).pop()
602 # check that it is a scalar
603 if const_tens.shape != []:
604 return op
605 const = const_tens.ops[0]
606 # check that it is a constant
607 if const.type != "Const":
608 return op
609 else:
610 return op
611
612 val = const.outputs[0].values
613 if val >= 0:
614 new_op = "LeakyRelu"
615 op.attrs["alpha"] = val
616 elif val == -1:
617 new_op = "Abs"
618 else:
619 return op
620
621 op.type = op.type.replace("Maximum", new_op)
622 op.name = op.name.replace("Maximum", new_op)
623 op.outputs[0].name = op.outputs[0].name.replace("Maximum", new_op)
624 op.inputs = [shared_in]
625 return op
626
627
Dwight Lidman42fed942020-05-29 09:37:03 +0200628def add_attrs_to_resizebilinear(op, arch):
Tim Hallc30f4952020-06-15 20:47:35 +0100629 if op.type == "ResizeBilinear" and op.run_on_npu:
Dwight Lidman42fed942020-05-29 09:37:03 +0200630 input_tensor = op.inputs[0]
631 upscaled_shape = [input_tensor.shape[1] * 2, input_tensor.shape[2] * 2]
632 out_shape = op.outputs[0].shape[1:3]
633 if not op.attrs["align_corners"] and out_shape == upscaled_shape:
634 # this means the output is supposed to be a x2 upscale,
635 # so we need to do SAME padding
636 op.attrs["padding"] = b"SAME"
637 elif op.attrs["align_corners"] and out_shape == [upscaled_shape[0] - 1, upscaled_shape[1] - 1]:
638 # here we can just run the avg pool without padding and
639 # produce a (M * 2 - 1, N * 2 - 1) sized output
640 op.attrs["padding"] = b"VALID"
641 else:
Charles Xu9a03fdf2020-07-02 15:12:40 +0200642 return op
Dwight Lidman42fed942020-05-29 09:37:03 +0200643 input_tensor.resampling_mode = resampling_mode.NEAREST
Tim Hallc30f4952020-06-15 20:47:35 +0100644 op.attrs.update({"strides": (1, 1, 1, 1), "ksize": (1, 2, 2, 1)})
Dwight Lidman42fed942020-05-29 09:37:03 +0200645 return op
646
647
Tim Hall79d07d22020-04-27 18:20:16 +0100648def supported_operator_check(op, arch):
649 op.run_on_npu = arch.supported_operators.is_operator_supported(op)
650 return op
651
652
653def optimise_graph_a(nng, arch, verbose_graph=False):
654 if verbose_graph:
655 nng.print_graph()
656
657 op_rewrite_list = [
658 # mark block type and check if the operations are supported
659 mark_npu_block_type,
Tim Hall4e127762020-05-15 16:05:49 +0100660 set_tensor_equivalence,
Tim Hall79d07d22020-04-27 18:20:16 +0100661 supported_operator_check,
662 # then do any rewrites of supported operators
663 convert_depthwise_to_conv,
Michael McGeagh8d939c02020-07-29 13:11:43 +0100664 convert_conv_to_fc,
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200665 convert_softmax,
Tim Hall79d07d22020-04-27 18:20:16 +0100666 fixup_fully_connected_input,
667 fixup_pack_input,
668 fixup_conv2d_backprop,
669 fixup_act_reorder,
Dwight Lidman42fed942020-05-29 09:37:03 +0200670 add_attrs_to_resizebilinear,
Tim Hall79d07d22020-04-27 18:20:16 +0100671 add_padding_fields,
672 mark_npu_block_type,
Charles Xu78792222020-05-13 10:15:26 +0200673 fixup_elementwise_with_scalars,
Jacob Bohline843d332020-06-23 12:12:56 +0200674 reorder_depthwise_weights,
Charles Xu9a03fdf2020-07-02 15:12:40 +0200675 fixup_resizebilinear,
Tim Hall79d07d22020-04-27 18:20:16 +0100676 # convert_mul_max_to_abs_or_lrelu # TODO: enable optimisation once quantisation issues are resolved
677 ]
678
679 for idx, sg in enumerate(nng.subgraphs):
680 # rewrite graph pass
681 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
Diego Russoea6111a2020-04-14 18:41:58 +0100682 sg, arch, [fixup_unpack_output], op_rewrite_list, rewrite_unsupported=False
Tim Hall79d07d22020-04-27 18:20:16 +0100683 )
684
685 for idx, sg in enumerate(nng.subgraphs):
686 # remove passthrough tensors
Diego Russoea6111a2020-04-14 18:41:58 +0100687 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(sg, arch, [remove_passthrough_tensor], [])
Tim Hall79d07d22020-04-27 18:20:16 +0100688
689 if verbose_graph:
690 nng.print_graph()
691 return nng
692
Diego Russoea6111a2020-04-14 18:41:58 +0100693
Tim Hall79d07d22020-04-27 18:20:16 +0100694def optimise_graph_b(nng, arch, verbose_graph=False):
695 if verbose_graph:
696 nng.print_graph()
697
698 for idx, sg in enumerate(nng.subgraphs):
699 # combined rewrite graph pass
Diego Russoea6111a2020-04-14 18:41:58 +0100700 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(sg, arch, [rewrite_concat, rewrite_split], [])
Tim Hall79d07d22020-04-27 18:20:16 +0100701
702 if verbose_graph:
703 nng.print_graph()
704 return nng