blob: a9d5cce517c24001fcd43640e307bdc095e6f5df [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
Charles Xu9a03fdf2020-07-02 15:12:40 +020031from .tensor import QuantizationParameters
Diego Russoe8a10452020-04-21 17:39:10 +010032from .tensor import Tensor
Tim Hall79d07d22020-04-27 18:20:16 +010033
34passthrough_nodes = set(("Identity",))
35
36
37def remove_passthrough_tensor(tens, arch):
38 if len(tens.ops) == 1 and tens.ops[0].type in passthrough_nodes:
39 assert len(tens.ops[0].inputs) == 1
40 tens = tens.ops[0].inputs[0]
41 return tens
42
43
44def rewrite_concat(tens, arch):
45 if len(tens.ops) == 1 and tens.ops[0].is_concat_op():
46 concat_op = tens.ops[0]
47 if tens != concat_op.outputs[0]:
48 return tens # don't attempt to rewrite the min/max outputs of QuantizedConcat
49
50 # Not supported so leave it and run on CPU
51 if not concat_op.run_on_npu:
52 return tens
53
54 inputs, axis = concat_op.get_concat_inputs_axis()
55
56 tens.ops = []
57 offset = 0
58 for idx, inp in enumerate(inputs):
59 new_op = Operation("ConcatSliceWrite", concat_op.name + str(idx))
60 new_op.inputs = [inp]
61 new_op.outputs = [tens]
62 new_op.attrs["concat_axis"] = axis
63 new_op.attrs["concat_start"] = offset
64 offset += inp.shape[axis]
65 new_op.attrs["concat_end"] = offset
66 new_op.run_on_npu = True
67 tens.ops.append(new_op)
68 assert tens.shape[axis] == offset
69
70 return tens
71
72
73def rewrite_split(tens, arch):
74
75 if len(tens.ops) == 1 and tens.ops[0].is_split_op():
76 split_op = tens.ops[0]
77
78 # Not supported so leave it and run on CPU
79 if not split_op.run_on_npu:
80 return tens
81
82 inp, outputs, axis, offset_start, offset_end = split_op.get_split_inputs_axis()
83
84 tens.ops = []
85 new_op = Operation("SplitSliceRead", split_op.name)
86 new_op.inputs = [inp]
87 new_op.outputs = [tens]
88
89 # For Split the offset cannot be extracted from the tensor so it has to
90 # be calculated from the index of the output tensor
Diego Russoea6111a2020-04-14 18:41:58 +010091 if axis is not None:
Tim Hall79d07d22020-04-27 18:20:16 +010092 # Get the start and end of the split
93 offset_start = [0] * len(tens.shape)
94 offset_end = [0] * len(tens.shape)
95 for out in outputs:
96 if out == tens:
97 break
98 offset_start[axis] += out.shape[axis]
99
100 offset_end[axis] = offset_start[axis] + tens.shape[axis]
101
102 new_op.attrs["split_start"] = offset_start
103 new_op.attrs["split_end"] = offset_end
104 new_op.run_on_npu = True
105 tens.ops.append(new_op)
106
107 return tens
108
109
110def needed_total_padding(input_size, stride, filter_size):
111 out_size = (input_size + stride - 1) // stride
112 needed_input = (out_size - 1) * stride + filter_size
113 total_padding = max(0, needed_input - input_size)
114 return total_padding
115
116
117def calc_padding_and_skirt(padding_type, kernel_size, stride, input_dims):
118 ypad = needed_total_padding(int(input_dims[1]), int(stride[1]), int(kernel_size[0]))
119 xpad = needed_total_padding(int(input_dims[2]), int(stride[2]), int(kernel_size[1]))
120 if padding_type == b"SAME":
121 left_pad = (xpad + 0) // 2
122 right_pad = (xpad + 1) // 2
123 top_pad = (ypad + 0) // 2
124 bottom_pad = (ypad + 1) // 2
125 elif padding_type == b"VALID":
126 left_pad = 0
127 right_pad = 0
128 top_pad = 0
129 bottom_pad = 0
130 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200131 raise UnsupportedFeatureError("Unknown padding {}".format(str(padding_type)))
Tim Hall79d07d22020-04-27 18:20:16 +0100132 padding = (top_pad, left_pad, bottom_pad, right_pad)
133 skirt = (top_pad, left_pad, ypad - top_pad, xpad - left_pad)
134 return padding, skirt
135
Tim Hallc30f4952020-06-15 20:47:35 +0100136
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200137def calc_upscaled_padding_and_skirt(padding_type, kernel_size, stride, input_dims, upscaling_factor):
138 kernel_height, kernel_width = kernel_size[0], kernel_size[1]
Jacob Bohlincf7da102020-05-20 09:03:40 +0200139 if padding_type == b"SAME":
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200140 ypad = needed_total_padding(int(input_dims[1]) * upscaling_factor, int(stride[1]), int(kernel_height))
141 xpad = needed_total_padding(int(input_dims[2]) * upscaling_factor, int(stride[2]), int(kernel_width))
142
143 right_pad = ((xpad + 1) // upscaling_factor) - 1
144 bottom_pad = ((ypad + 1) // upscaling_factor) - 1
145 left_pad = max(kernel_width - 1 - right_pad, 0)
146 top_pad = max(kernel_height - 1 - bottom_pad, 0)
147
Jacob Bohlincf7da102020-05-20 09:03:40 +0200148 elif padding_type == b"VALID":
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200149 right_pad = max(kernel_width - 2, 0)
150 bottom_pad = max(kernel_height - 2, 0)
151 left_pad = kernel_width - 1
152 top_pad = kernel_height - 1
Jacob Bohlincf7da102020-05-20 09:03:40 +0200153 else:
154 assert 0, "Unknown padding"
155
156 padding = (top_pad, left_pad, bottom_pad, right_pad)
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200157 skirt = padding
Jacob Bohlincf7da102020-05-20 09:03:40 +0200158 return padding, skirt
159
Tim Hall79d07d22020-04-27 18:20:16 +0100160
161def fixup_conv2d_backprop(op, arch):
162 if op.type == "Conv2DBackpropInput":
163 # flip the inputs
164 op.inputs[0], op.inputs[2] = op.inputs[2], op.inputs[0]
Jacob Bohlincf7da102020-05-20 09:03:40 +0200165 op.type = "Conv2DBackpropInputSwitchedBias"
166 weight_shape = op.inputs[1].shape
167 weight_sets = weight_shape[3]
168
169 if len(op.inputs) < 4:
170 # Add bias/scale tensor filled with zeros
171 scale_op = Operation("Const", op.name + "_bias")
172 scale_tens = Tensor([weight_sets], DataType.int32, op.name + "_bias_tens")
173 scale_tens.values = [0] * weight_sets
174 scale_tens.quant_values = [0] * weight_sets
175 scale_tens.ops = [scale_op]
176 scale_op.outputs = [scale_tens]
177 scale_tens.consumer_list = [op]
178 op.inputs.append(scale_tens)
179
180 # Update strides
Tim Hallc30f4952020-06-15 20:47:35 +0100181 op.attrs.update({"stride_w": 1, "stride_h": 1, "strides": (1, 1, 1, 1)})
Tim Hall79d07d22020-04-27 18:20:16 +0100182
183 return op
184
185
Charles Xu9a03fdf2020-07-02 15:12:40 +0200186# Convert the op to an elementwise add
187def convert_resizebilinear_1x1_to_add(op):
188 op.type = "AddAct"
189 op.name = op.name + "_add"
190 op.attrs.update({"npu_block_type": NpuBlockType.ElementWise})
191 op.attrs["resizebilinear"] = True
192 # Create an input tensor filled with zeros
193 shape = op.outputs[0].shape
194 tens = Tensor(shape, op.inputs[0].dtype, op.inputs[1].name + "_add")
195 tens.values = np.zeros(shape)
196 tens.quant_values = np.zeros(shape, np.uint8)
197 tens.quantization = QuantizationParameters(0.0, 255.0)
198 tens.quantization.scale_f32 = 1.0
199 tens.quantization.zero_point = 0
200 tens.consumer_list = [op]
201 tens_op = op.inputs[1].ops[0]
202 tens_op.outputs = [tens]
203 tens.ops = [tens_op]
204 # Set the add inputs
205 op.inputs[1] = op.inputs[0]
206 op.inputs[0] = tens
207
208 return op
209
210
211def fixup_resizebilinear(op, arch):
212 if op.type == "ResizeBilinear":
213 if op.inputs[0].shape[1] == 1 and op.inputs[0].shape[2] == 1:
214 convert_resizebilinear_1x1_to_add(op)
Charles Xu36ffaf32020-08-05 15:40:44 +0200215 elif op.inputs[0].shape == op.outputs[0].shape:
216 # Bypass nop resizebilinear
217 op.inputs = op.inputs[:1]
218 op.type = "Identity"
Charles Xu9a03fdf2020-07-02 15:12:40 +0200219
220 return op
221
222
Tim Hall79d07d22020-04-27 18:20:16 +0100223def fixup_fully_connected_input(op, arch):
224 if op.type == "FullyConnectedAct":
225 inp = op.inputs[0]
226 weights = op.inputs[1]
227
228 n_in_elems = weights.shape[-2]
229 elms = inp.elements()
230 batch_size = elms // n_in_elems
231 assert batch_size * n_in_elems == elms
232
233 desired_shape = [batch_size, n_in_elems]
234 if inp.shape != desired_shape:
235 # mismatch, insert a reshape to fix this.
236 reshape_name = op.name + "_reshape"
237 new_shape_tens = Tensor([1], DataType.int32, reshape_name + "_shape")
238 new_shape_tens.values = np.array(desired_shape)
239 new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
240 new_shape_tens.ops = [new_shape_tens_const]
241 new_shape_tens_const.outputs = [new_shape_tens]
242
243 reshape_op = Operation("Reshape", reshape_name)
244 reshape_op.inputs = [inp, new_shape_tens]
245 reshape_op.attrs["new_shape"] = desired_shape
246 reshape_out = inp.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100247 reshape_out.set_all_shapes(desired_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100248 reshape_out.ops = [reshape_op]
249 reshape_op.outputs = [reshape_out]
250
251 op.inputs[0] = reshape_out
252
253 return op
254
255
256def fixup_pack_input(op, arch):
257 if op.type == "Pack":
258 # Pack is also referred to as Stack
259 # Requires the rewrite_concat function to be called on the op afterwards
260 axis = int(op.attrs["axis"])
261 desired_shape = op.inputs[0].shape[:axis] + [1] + op.inputs[0].shape[axis:]
262
263 # Construct 1 shape tensor to be used by all inserted reshape ops
264 new_shape_name = op.name + "_reshape_shape"
265 new_shape_tens = Tensor([1], DataType.int32, new_shape_name)
266 new_shape_tens.values = np.array(desired_shape)
267 new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
268 new_shape_tens.ops = [new_shape_tens_const]
269 new_shape_tens_const.outputs = [new_shape_tens]
270
271 for idx, inp in enumerate(op.inputs):
272 reshape_name = op.name + str(idx) + "_reshape"
273 reshape_op = Operation("Reshape", reshape_name)
274 reshape_op.inputs = [inp, new_shape_tens]
275 reshape_op.attrs["new_shape"] = desired_shape
276 reshape_out = inp.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100277 reshape_out.set_all_shapes(desired_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100278 reshape_out.ops = [reshape_op]
279 reshape_op.outputs = [reshape_out]
280
281 op.inputs[idx] = reshape_out
282
283 op.type = "PackReshaped"
284
285 return op
286
287
288def fixup_unpack_output(tens, arch):
289 op = tens.ops[0]
290 if op.type in set(("Unpack", "StridedSlice")):
291 # Unpack is also referred to as Unstack
292 # Requires the rewrite_split function to be called on the op afterwards
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200293
294 reshape_input_shape = tens.shape
Tim Hall79d07d22020-04-27 18:20:16 +0100295 if op.type == "StridedSlice":
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200296 new_axis_mask = op.attrs["new_axis_mask"]
Tim Hall79d07d22020-04-27 18:20:16 +0100297 shrink_axis_mask = op.attrs["shrink_axis_mask"]
Louis Verhaard7db78962020-05-25 15:05:26 +0200298 ellipsis_mask = op.attrs["ellipsis_mask"]
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200299
300 if (new_axis_mask != 0 and shrink_axis_mask != 0) or ellipsis_mask != 0:
301 # Not supported, will be put on CPU
302 return tens
303 if shrink_axis_mask == 0 and new_axis_mask == 0:
Tim Hall79d07d22020-04-27 18:20:16 +0100304 # Equal Rank StridedSlice, no need to insert reshape
305 return tens
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200306 elif shrink_axis_mask != 0:
307 n = 0
308 axis = 0
309 while shrink_axis_mask:
310 prev_mask = shrink_axis_mask
311 n += 1
312 shrink_axis_mask &= shrink_axis_mask - 1
313 axis = int(math.log2(prev_mask - shrink_axis_mask))
314 reshape_input_shape = reshape_input_shape[:axis] + [1] + reshape_input_shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100315
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200316 assert len(tens.shape) == (len(op.inputs[0].shape) - n)
317 op.attrs["shrink_axis_mask"] = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100318
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200319 elif new_axis_mask != 0:
320 n = 0
321 axis = 0
322 while new_axis_mask:
323 prev_mask = new_axis_mask
324 n += 1
325 new_axis_mask &= new_axis_mask - 1
326 axis = int(math.log2(prev_mask - new_axis_mask))
Louis Verhaard7db78962020-05-25 15:05:26 +0200327 reshape_input_shape = reshape_input_shape[:axis] + reshape_input_shape[(axis + 1) :]
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200328 new_axis_mask >>= 1
329
330 assert len(tens.shape) == (len(op.inputs[0].shape) + n)
331 op.attrs["new_axis_mask"] = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100332 else:
333 axis = int(op.attrs["axis"])
334 op.type = "UnpackReshaped"
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200335 reshape_input_shape = tens.shape[:axis] + [1] + tens.shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100336
337 # Construct 1 shape tensor to be used by all inserted reshape ops
338 new_shape_name = op.name + "_reshape_shape"
339 new_shape_tens = Tensor([1], DataType.int32, new_shape_name)
340 new_shape_tens.values = np.array(tens.shape)
341 new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
342 new_shape_tens.ops = [new_shape_tens_const]
343 new_shape_tens_const.outputs = [new_shape_tens]
344
345 for idx, out_tens in enumerate(op.outputs):
346 reshape_name = op.name + str(idx) + "_reshape"
347 reshape_op = Operation("Reshape", reshape_name)
348 reshape_op.outputs = [out_tens]
349 reshape_in = out_tens.clone("_reshaped")
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100350 reshape_in.set_all_shapes(reshape_input_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100351 reshape_in.ops = [op]
352 out_tens.ops = [reshape_op]
353 reshape_op.inputs = [reshape_in, new_shape_tens]
354
355 op.outputs[idx] = reshape_in
356
357 return tens
358
359
360def add_padding_fields(op, arch):
361 if "padding" in op.attrs:
362 if "Conv" in op.type:
363 kernel_size = op.inputs[1].shape[:2]
364 input_shape = op.inputs[0].shape
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200365 elif "Pool" in op.type or op.type in ("ResizeBilinear", "ReduceSum"):
Tim Hall79d07d22020-04-27 18:20:16 +0100366 kernel_size = op.attrs["ksize"][1:3]
367 input_shape = op.inputs[0].shape
368 elif op.type == "ExtractImagePatches":
369 kernel_size = op.attrs["ksizes"][1:3]
370 input_shape = op.inputs[0].shape
371 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200372 raise UnsupportedFeatureError("Unknown operation that uses padding: {}".format(op.type))
Tim Hall79d07d22020-04-27 18:20:16 +0100373
Jacob Bohlincf7da102020-05-20 09:03:40 +0200374 if op.type == "Conv2DBackpropInputSwitchedBias":
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200375 upscaling_factor = op.outputs[0].shape[1] // input_shape[1]
Tim Hallc30f4952020-06-15 20:47:35 +0100376 padding, skirt = calc_upscaled_padding_and_skirt(
Jacob Bohlin9b64ba02020-07-07 17:15:22 +0200377 op.attrs["padding"], kernel_size, op.attrs["strides"], input_shape, upscaling_factor
Tim Hallc30f4952020-06-15 20:47:35 +0100378 )
Jacob Bohlincf7da102020-05-20 09:03:40 +0200379 else:
380 dilation_h, dilation_w = op.get_dilation_h_w()
381 dilated_kernel_size = [dilation_h * (kernel_size[0] - 1) + 1, dilation_w * (kernel_size[1] - 1) + 1]
Tim Hallc30f4952020-06-15 20:47:35 +0100382 padding, skirt = calc_padding_and_skirt(
383 op.attrs["padding"], dilated_kernel_size, op.attrs["strides"], input_shape
384 )
Jacob Bohlincf7da102020-05-20 09:03:40 +0200385
Tim Hall79d07d22020-04-27 18:20:16 +0100386 op.attrs["explicit_padding"] = padding
387 op.attrs["skirt"] = skirt
Jacob Bohlincf7da102020-05-20 09:03:40 +0200388
Tim Hall79d07d22020-04-27 18:20:16 +0100389 return op
390
391
Jacob Bohlincf7da102020-05-20 09:03:40 +0200392conv_op = set(("Conv2D", "QuantizedConv2D", "Conv2DBackpropInputSwitchedBias", "Conv2DBiasAct"))
Tim Hall79d07d22020-04-27 18:20:16 +0100393fc_op = set(
394 (
395 "MatMul",
396 "QuantizedMatMul",
397 "BlockLSTM",
398 "RnnAct",
399 "UnidirectionalSequenceRnnAct",
400 "BidirectionalSequenceRnnAct",
401 "LstmAct",
402 "UnidirectionalSequenceLstmAct",
403 "BidirectionalSequenceLstmAct",
404 "FullyConnectedAct",
405 )
406)
407depthwise_op = set(("DepthwiseConv2dNative", "DepthwiseConv2dBiasAct",))
Louis Verhaard7db78962020-05-25 15:05:26 +0200408pool_op = set(
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200409 ("AvgPool", "MaxPool", "QuantizedAvgPool", "QuantizedMaxPool", "AvgPoolAct", "MaxPoolAct", "ResizeBilinear")
Louis Verhaard7db78962020-05-25 15:05:26 +0200410)
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200411reduce_sum_ops = set(("ReduceSum",))
412elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum", "LeakyRelu", "Abs", "CLZ", "SHL", "SHR"))
Charles Xu78792222020-05-13 10:15:26 +0200413binary_elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum"))
Tim Hall79d07d22020-04-27 18:20:16 +0100414activation_ops = set(("Relu", "Relu6", "ReluN1To1", "Sigmoid", "Tanh"))
415memory_only_ops = set(("Reshape",))
416
Diego Russoea6111a2020-04-14 18:41:58 +0100417
Tim Hall79d07d22020-04-27 18:20:16 +0100418# Check if the op can be reordered
419def get_prepend_op(op):
420 inp = op.inputs[0]
421 # The op should be reordered between prev_op and prep_op
422 prev_op = inp.ops[-1]
423 prep_op = None
424 while prev_op.type in memory_only_ops and len(prev_op.outputs) == 1 and len(prev_op.outputs[0].consumers()) == 1:
425 prep_op = prev_op
426 inp = prev_op.inputs[0]
427 prev_op = inp.ops[-1]
Diego Russoea6111a2020-04-14 18:41:58 +0100428 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 +0100429 return prep_op
430
431 return None
432
433
434def mark_npu_block_type(op, arch):
435 npu_block_type = NpuBlockType.Default
436 if op.type in conv_op:
437 npu_block_type = NpuBlockType.ConvolutionMxN
438 elif op.type in fc_op:
439 npu_block_type = NpuBlockType.VectorProduct
440 elif op.type in depthwise_op:
441 npu_block_type = NpuBlockType.ConvolutionDepthWise
442 elif op.type in pool_op:
443 npu_block_type = NpuBlockType.Pooling
444 elif op.type in elementwise_op:
445 npu_block_type = NpuBlockType.ElementWise
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200446 elif op.type in reduce_sum_ops:
447 npu_block_type = NpuBlockType.ReduceSum
Tim Hall79d07d22020-04-27 18:20:16 +0100448
449 op.attrs["npu_block_type"] = npu_block_type
450 return op
451
452
453def convert_depthwise_to_conv(op, arch):
454 # Depthwise is equivalent to a single conv2d if the ifm depth is 1 and
455 # the ofm depth equals the depth multipler.
456 # If those conditions are true, then we can perform a simple
457 # switch of the operator type (and weight order)
458
459 if ("DepthwiseConv2d" in op.type) and (op.attrs["depth_multiplier"] != 1):
460 ifm_tensor = op.inputs[0]
461 weight_tensor = op.inputs[1]
462 ofm_tensor = op.outputs[0]
463 if (ifm_tensor.shape[3] == 1) and (ofm_tensor.shape[3] == op.attrs["depth_multiplier"]):
464 # Change op type to Conv2d
465 op.type = op.type.replace("DepthwiseConv2d", "Conv2D")
466 del op.attrs["channel_multiplier"]
467 del op.attrs["depth_multiplier"]
468
469 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100470 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Tim Hall79d07d22020-04-27 18:20:16 +0100471 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200472 raise UnsupportedFeatureError(
473 "Unsupported DepthwiseConv2d with depth_multiplier = {}, ifm channels = {}, ofm channels = {}".format(
Tim Hall79d07d22020-04-27 18:20:16 +0100474 op.attrs["depth_multiplier"], ifm_tensor.shape[3], ofm_tensor.shape[3]
475 )
476 )
Tim Hall79d07d22020-04-27 18:20:16 +0100477 return op
478
479
Jacob Bohline843d332020-06-23 12:12:56 +0200480def reorder_depthwise_weights(op, arch):
481 if "DepthwiseConv2d" in op.type:
482 weight_tensor = op.inputs[1]
483 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100484 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
Jacob Bohline843d332020-06-23 12:12:56 +0200485 weight_tensor.weight_transpose_depthwise = True
486
487 return op
488
489
Michael McGeagh8d939c02020-07-29 13:11:43 +0100490def convert_conv_to_fc(op, arch):
491 # Conv 1x1 can be equivalent to Fully Connected.
492 # By representing certain convs as fully connected layers, Vela can better determine wether or not to use
493 # caching/double buffering for the weights.
494 # (Weights dont need to be reloaded for convs when IFM H and W are 1)
495 if op.type == "Conv2DBiasAct":
496 _, h, w, _ = op.inputs[0].shape
497 kh, kw, _, _ = op.inputs[1].shape
498 if h == 1 and w == 1 and kh == 1 and kw == 1:
499 # Overwrite this op as a Fully Connected Op
500 op.name += "_fc"
501 op.type = "FullyConnectedAct"
502 faf = op.attrs.get("fused_activation_function", None)
503 op.attrs = {
504 "fused_activation_function": faf,
505 "weights_format": 0,
506 "npu_block_type": NpuBlockType.VectorProduct,
507 }
508 # Reshape Weights to be 2D. HWIO becomes just IO (as H and W are 1, they can just be dropped)
509 weight_tensor = op.inputs[1]
510 weight_tensor.quant_values = weight_tensor.quant_values.squeeze(axis=(0, 1))
511 weight_tensor.set_all_shapes(list(weight_tensor.quant_values.shape))
512 # The output from a fully connected is expected to be 2D so we need to add a reshape layer to convert it
513 # back to 4D afterwards as the next layer is expecting that shape
514 orig_ofm_tensor = op.outputs[0]
515 # 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})
516 fc_ofm_tensor = orig_ofm_tensor.clone("_fc")
517 fc_ofm_tensor.set_all_shapes([1, fc_ofm_tensor.shape[-1]])
518 fc_ofm_tensor.ops = [op]
519 # Add a reshape after the new OFM to convert it back to the original 4D shape
520 reshape_name = op.name + "_reshape_post"
521 new_shape_tens = Tensor([1], DataType.int32, reshape_name + "_shape")
522 new_shape_tens.values = np.array(orig_ofm_tensor.shape)
523 new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
524 new_shape_tens.ops = [new_shape_tens_const]
525 new_shape_tens_const.outputs = [new_shape_tens]
526 reshape_op = Operation("Reshape", reshape_name)
527 reshape_op.inputs = [fc_ofm_tensor, new_shape_tens]
528 reshape_op.attrs["new_shape"] = orig_ofm_tensor.shape
529 orig_ofm_tensor.ops = [reshape_op]
530 reshape_op.outputs = [orig_ofm_tensor]
531 # Replace this ops OFM to point to the 2D tensor
532 op.outputs[0] = fc_ofm_tensor
533 return op
534
535
Tim Hall79d07d22020-04-27 18:20:16 +0100536# Reorder activation op if it's after the memory only operations
537def fixup_act_reorder(op, arch):
538 if op.type in activation_ops:
539 prep_op = get_prepend_op(op)
Diego Russoea6111a2020-04-14 18:41:58 +0100540 if prep_op is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100541 act_op = op.clone("_reordered")
542 act_op.inputs = [prep_op.inputs[0]]
543 act_op_out = act_op.inputs[0].clone("_acted")
544 act_op_out.quantization = op.outputs[0].quantization.clone()
545 act_op_out.ops = [act_op]
546 act_op.outputs = [act_op_out]
547 prep_op.inputs[0] = act_op_out
548 prep_op.outputs[0].quantization = act_op_out.quantization.clone()
549
550 # Mark the op so that it will be removed as passthrough later on
551 op.type = "Identity"
552 return op
553
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200554
Charles Xu78792222020-05-13 10:15:26 +0200555def fixup_elementwise_with_scalars(op, arch):
556 if op.type in binary_elementwise_op:
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200557 ifm_tensor, ifm2_tensor, _, _ = op.get_ifm_ifm2_weights_ofm()
Charles Xu78792222020-05-13 10:15:26 +0200558 if ifm2_tensor.shape != [] and ifm_tensor.shape != []:
559 diff = len(ifm_tensor.shape) - len(ifm2_tensor.shape)
560 if diff > 0:
561 ifm2_tensor.shape = full_shape(len(ifm_tensor.shape), ifm2_tensor.shape, 1)
562 elif diff < 0:
563 ifm_tensor.shape = full_shape(len(ifm2_tensor.shape), ifm_tensor.shape, 1)
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200564 elif ifm_tensor.shape == [] and ifm_tensor.quant_values is None:
565 # IFM is marked as a scalar, but is a result of an operation; change it to a shape of size 1
566 ifm_tensor.shape = len(ifm2_tensor.shape) * [1]
567 ifm_tensor.storage_shape = ifm_tensor.shape
568 elif ifm2_tensor.shape == [] and ifm2_tensor.quant_values is None:
569 # IFM2 is marked as a scalar, but is a result of an operation; change it to a shape of size 1
570 ifm2_tensor.shape = len(ifm_tensor.shape) * [1]
571 ifm2_tensor.storage_shape = ifm2_tensor.shape
Charles Xu78792222020-05-13 10:15:26 +0200572 return op
Tim Hall79d07d22020-04-27 18:20:16 +0100573
Louis Verhaarde0ef2732020-06-03 08:56:44 +0200574
Tim Hall4e127762020-05-15 16:05:49 +0100575# Set input/output tensor equivalence to the same id for memory operations
576def set_tensor_equivalence(op, arch):
577 if op.type == "Reshape":
578 eid = op.outputs[0].equivalence_id
579 for inp in op.inputs:
580 inp.equivalence_id = eid
581 return op
582
583
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200584def convert_softmax(op, arch):
585 if op.type == "Softmax" and op.run_on_npu:
586 softmax = SoftMax(op)
587 op = softmax.get_graph()
588 return op
589
590
Tim Hall79d07d22020-04-27 18:20:16 +0100591def convert_mul_max_to_abs_or_lrelu(op, arch):
Diego Russoea6111a2020-04-14 18:41:58 +0100592 r"""Whenever there is a subgraph with this topology:
Tim Hall79d07d22020-04-27 18:20:16 +0100593
594 Input X For X = -1 or X > 0
595 | \ / This subgraph can be replaced with either
596 | Mul an Abs (if X = -1) or a LeakyReLU (if X > 0)
597 | /
598 Max
599 """
600
601 if op.type == "Maximum":
602 # finds the Mul input(s) to the Max
603 muls = [i for i in op.inputs if i.ops[0].type == "MulAct"]
604 if len(muls) == 1:
605 mul = muls[0].ops[0]
606 elif len(muls) == 2:
607 # In the case both inputs are Muls, find the one with the same input as the Max
608 mul = [m for m in muls if len(set(op.inputs + m.ops[0].inputs)) == 1][0].ops[0]
609 else:
610 # No Mul inputs
611 return op
612
613 # make sure the Mul doesn't have any other consumers
614 if len(mul.outputs[0].consumers()) != 1:
615 return op
616 # make sure the Mul doesn't have a faf
617 if mul.attrs["fused_activation_function"]:
618 return op
619
620 # finds the branched input that goes to both the Max and the Mul
621 shared = set(op.inputs) & set(mul.inputs)
622 if len(shared) == 1:
623 shared_in = shared.pop()
624 # find the constant scalar input to the Mul
625 const_tens = (set(mul.inputs) - {shared_in}).pop()
626 # check that it is a scalar
627 if const_tens.shape != []:
628 return op
629 const = const_tens.ops[0]
630 # check that it is a constant
631 if const.type != "Const":
632 return op
633 else:
634 return op
635
636 val = const.outputs[0].values
637 if val >= 0:
638 new_op = "LeakyRelu"
639 op.attrs["alpha"] = val
640 elif val == -1:
641 new_op = "Abs"
642 else:
643 return op
644
645 op.type = op.type.replace("Maximum", new_op)
646 op.name = op.name.replace("Maximum", new_op)
647 op.outputs[0].name = op.outputs[0].name.replace("Maximum", new_op)
648 op.inputs = [shared_in]
649 return op
650
651
Dwight Lidman42fed942020-05-29 09:37:03 +0200652def add_attrs_to_resizebilinear(op, arch):
Tim Hallc30f4952020-06-15 20:47:35 +0100653 if op.type == "ResizeBilinear" and op.run_on_npu:
Dwight Lidman42fed942020-05-29 09:37:03 +0200654 input_tensor = op.inputs[0]
655 upscaled_shape = [input_tensor.shape[1] * 2, input_tensor.shape[2] * 2]
656 out_shape = op.outputs[0].shape[1:3]
657 if not op.attrs["align_corners"] and out_shape == upscaled_shape:
658 # this means the output is supposed to be a x2 upscale,
659 # so we need to do SAME padding
660 op.attrs["padding"] = b"SAME"
661 elif op.attrs["align_corners"] and out_shape == [upscaled_shape[0] - 1, upscaled_shape[1] - 1]:
662 # here we can just run the avg pool without padding and
663 # produce a (M * 2 - 1, N * 2 - 1) sized output
664 op.attrs["padding"] = b"VALID"
665 else:
Charles Xu9a03fdf2020-07-02 15:12:40 +0200666 return op
Dwight Lidman42fed942020-05-29 09:37:03 +0200667 input_tensor.resampling_mode = resampling_mode.NEAREST
Tim Hallc30f4952020-06-15 20:47:35 +0100668 op.attrs.update({"strides": (1, 1, 1, 1), "ksize": (1, 2, 2, 1)})
Dwight Lidman42fed942020-05-29 09:37:03 +0200669 return op
670
671
Tim Hall79d07d22020-04-27 18:20:16 +0100672def supported_operator_check(op, arch):
673 op.run_on_npu = arch.supported_operators.is_operator_supported(op)
674 return op
675
676
677def optimise_graph_a(nng, arch, verbose_graph=False):
678 if verbose_graph:
679 nng.print_graph()
680
681 op_rewrite_list = [
682 # mark block type and check if the operations are supported
683 mark_npu_block_type,
Tim Hall4e127762020-05-15 16:05:49 +0100684 set_tensor_equivalence,
Tim Hall79d07d22020-04-27 18:20:16 +0100685 supported_operator_check,
686 # then do any rewrites of supported operators
687 convert_depthwise_to_conv,
Michael McGeagh8d939c02020-07-29 13:11:43 +0100688 convert_conv_to_fc,
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200689 convert_softmax,
Tim Hall79d07d22020-04-27 18:20:16 +0100690 fixup_fully_connected_input,
691 fixup_pack_input,
692 fixup_conv2d_backprop,
693 fixup_act_reorder,
Dwight Lidman42fed942020-05-29 09:37:03 +0200694 add_attrs_to_resizebilinear,
Tim Hall79d07d22020-04-27 18:20:16 +0100695 add_padding_fields,
696 mark_npu_block_type,
Charles Xu78792222020-05-13 10:15:26 +0200697 fixup_elementwise_with_scalars,
Jacob Bohline843d332020-06-23 12:12:56 +0200698 reorder_depthwise_weights,
Charles Xu9a03fdf2020-07-02 15:12:40 +0200699 fixup_resizebilinear,
Tim Hall79d07d22020-04-27 18:20:16 +0100700 # convert_mul_max_to_abs_or_lrelu # TODO: enable optimisation once quantisation issues are resolved
701 ]
702
703 for idx, sg in enumerate(nng.subgraphs):
704 # rewrite graph pass
705 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
Diego Russoea6111a2020-04-14 18:41:58 +0100706 sg, arch, [fixup_unpack_output], op_rewrite_list, rewrite_unsupported=False
Tim Hall79d07d22020-04-27 18:20:16 +0100707 )
708
709 for idx, sg in enumerate(nng.subgraphs):
710 # remove passthrough tensors
Diego Russoea6111a2020-04-14 18:41:58 +0100711 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(sg, arch, [remove_passthrough_tensor], [])
Tim Hall79d07d22020-04-27 18:20:16 +0100712
713 if verbose_graph:
714 nng.print_graph()
715 return nng
716
Diego Russoea6111a2020-04-14 18:41:58 +0100717
Tim Hall79d07d22020-04-27 18:20:16 +0100718def optimise_graph_b(nng, arch, verbose_graph=False):
719 if verbose_graph:
720 nng.print_graph()
721
722 for idx, sg in enumerate(nng.subgraphs):
723 # combined rewrite graph pass
Diego Russoea6111a2020-04-14 18:41:58 +0100724 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(sg, arch, [rewrite_concat, rewrite_split], [])
Tim Hall79d07d22020-04-27 18:20:16 +0100725
726 if verbose_graph:
727 nng.print_graph()
728 return nng