blob: 72bb486c0215573a1d1eb35608dcc5edd21ab5fa [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
Diego Russoe8a10452020-04-21 17:39:10 +010026from .operation import NpuBlockType
27from .operation import Operation
28from .tensor import Tensor
Charles Xu78792222020-05-13 10:15:26 +020029from .numeric_util import full_shape
Tim Hall79d07d22020-04-27 18:20:16 +010030
31passthrough_nodes = set(("Identity",))
32
33
34def remove_passthrough_tensor(tens, arch):
35 if len(tens.ops) == 1 and tens.ops[0].type in passthrough_nodes:
36 assert len(tens.ops[0].inputs) == 1
37 tens = tens.ops[0].inputs[0]
38 return tens
39
40
41def rewrite_concat(tens, arch):
42 if len(tens.ops) == 1 and tens.ops[0].is_concat_op():
43 concat_op = tens.ops[0]
44 if tens != concat_op.outputs[0]:
45 return tens # don't attempt to rewrite the min/max outputs of QuantizedConcat
46
47 # Not supported so leave it and run on CPU
48 if not concat_op.run_on_npu:
49 return tens
50
51 inputs, axis = concat_op.get_concat_inputs_axis()
52
53 tens.ops = []
54 offset = 0
55 for idx, inp in enumerate(inputs):
56 new_op = Operation("ConcatSliceWrite", concat_op.name + str(idx))
57 new_op.inputs = [inp]
58 new_op.outputs = [tens]
59 new_op.attrs["concat_axis"] = axis
60 new_op.attrs["concat_start"] = offset
61 offset += inp.shape[axis]
62 new_op.attrs["concat_end"] = offset
63 new_op.run_on_npu = True
64 tens.ops.append(new_op)
65 assert tens.shape[axis] == offset
66
67 return tens
68
69
70def rewrite_split(tens, arch):
71
72 if len(tens.ops) == 1 and tens.ops[0].is_split_op():
73 split_op = tens.ops[0]
74
75 # Not supported so leave it and run on CPU
76 if not split_op.run_on_npu:
77 return tens
78
79 inp, outputs, axis, offset_start, offset_end = split_op.get_split_inputs_axis()
80
81 tens.ops = []
82 new_op = Operation("SplitSliceRead", split_op.name)
83 new_op.inputs = [inp]
84 new_op.outputs = [tens]
85
86 # For Split the offset cannot be extracted from the tensor so it has to
87 # be calculated from the index of the output tensor
Diego Russoea6111a2020-04-14 18:41:58 +010088 if axis is not None:
Tim Hall79d07d22020-04-27 18:20:16 +010089 # Get the start and end of the split
90 offset_start = [0] * len(tens.shape)
91 offset_end = [0] * len(tens.shape)
92 for out in outputs:
93 if out == tens:
94 break
95 offset_start[axis] += out.shape[axis]
96
97 offset_end[axis] = offset_start[axis] + tens.shape[axis]
98
99 new_op.attrs["split_start"] = offset_start
100 new_op.attrs["split_end"] = offset_end
101 new_op.run_on_npu = True
102 tens.ops.append(new_op)
103
104 return tens
105
106
107def needed_total_padding(input_size, stride, filter_size):
108 out_size = (input_size + stride - 1) // stride
109 needed_input = (out_size - 1) * stride + filter_size
110 total_padding = max(0, needed_input - input_size)
111 return total_padding
112
113
114def calc_padding_and_skirt(padding_type, kernel_size, stride, input_dims):
115 ypad = needed_total_padding(int(input_dims[1]), int(stride[1]), int(kernel_size[0]))
116 xpad = needed_total_padding(int(input_dims[2]), int(stride[2]), int(kernel_size[1]))
117 if padding_type == b"SAME":
118 left_pad = (xpad + 0) // 2
119 right_pad = (xpad + 1) // 2
120 top_pad = (ypad + 0) // 2
121 bottom_pad = (ypad + 1) // 2
122 elif padding_type == b"VALID":
123 left_pad = 0
124 right_pad = 0
125 top_pad = 0
126 bottom_pad = 0
127 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200128 raise UnsupportedFeatureError("Unknown padding {}".format(str(padding_type)))
Tim Hall79d07d22020-04-27 18:20:16 +0100129 padding = (top_pad, left_pad, bottom_pad, right_pad)
130 skirt = (top_pad, left_pad, ypad - top_pad, xpad - left_pad)
131 return padding, skirt
132
133
134def fixup_conv2d_backprop(op, arch):
135 if op.type == "Conv2DBackpropInput":
136 # flip the inputs
137 op.inputs[0], op.inputs[2] = op.inputs[2], op.inputs[0]
138 op.type = "Conv2DBackpropInputSwitched"
139
140 return op
141
142
143def fixup_fully_connected_input(op, arch):
144 if op.type == "FullyConnectedAct":
145 inp = op.inputs[0]
146 weights = op.inputs[1]
147
148 n_in_elems = weights.shape[-2]
149 elms = inp.elements()
150 batch_size = elms // n_in_elems
151 assert batch_size * n_in_elems == elms
152
153 desired_shape = [batch_size, n_in_elems]
154 if inp.shape != desired_shape:
155 # mismatch, insert a reshape to fix this.
156 reshape_name = op.name + "_reshape"
157 new_shape_tens = Tensor([1], DataType.int32, reshape_name + "_shape")
158 new_shape_tens.values = np.array(desired_shape)
159 new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
160 new_shape_tens.ops = [new_shape_tens_const]
161 new_shape_tens_const.outputs = [new_shape_tens]
162
163 reshape_op = Operation("Reshape", reshape_name)
164 reshape_op.inputs = [inp, new_shape_tens]
165 reshape_op.attrs["new_shape"] = desired_shape
166 reshape_out = inp.clone("_reshaped")
167 reshape_out.shape = reshape_out.storage_shape = reshape_out.bandwidth_shape = desired_shape
168 reshape_out.ops = [reshape_op]
169 reshape_op.outputs = [reshape_out]
170
171 op.inputs[0] = reshape_out
172
173 return op
174
175
176def fixup_pack_input(op, arch):
177 if op.type == "Pack":
178 # Pack is also referred to as Stack
179 # Requires the rewrite_concat function to be called on the op afterwards
180 axis = int(op.attrs["axis"])
181 desired_shape = op.inputs[0].shape[:axis] + [1] + op.inputs[0].shape[axis:]
182
183 # Construct 1 shape tensor to be used by all inserted reshape ops
184 new_shape_name = op.name + "_reshape_shape"
185 new_shape_tens = Tensor([1], DataType.int32, new_shape_name)
186 new_shape_tens.values = np.array(desired_shape)
187 new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
188 new_shape_tens.ops = [new_shape_tens_const]
189 new_shape_tens_const.outputs = [new_shape_tens]
190
191 for idx, inp in enumerate(op.inputs):
192 reshape_name = op.name + str(idx) + "_reshape"
193 reshape_op = Operation("Reshape", reshape_name)
194 reshape_op.inputs = [inp, new_shape_tens]
195 reshape_op.attrs["new_shape"] = desired_shape
196 reshape_out = inp.clone("_reshaped")
197 reshape_out.shape = reshape_out.storage_shape = reshape_out.bandwidth_shape = desired_shape
198 reshape_out.ops = [reshape_op]
199 reshape_op.outputs = [reshape_out]
200
201 op.inputs[idx] = reshape_out
202
203 op.type = "PackReshaped"
204
205 return op
206
207
208def fixup_unpack_output(tens, arch):
209 op = tens.ops[0]
210 if op.type in set(("Unpack", "StridedSlice")):
211 # Unpack is also referred to as Unstack
212 # Requires the rewrite_split function to be called on the op afterwards
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200213
214 reshape_input_shape = tens.shape
Tim Hall79d07d22020-04-27 18:20:16 +0100215 if op.type == "StridedSlice":
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200216 new_axis_mask = op.attrs["new_axis_mask"]
Tim Hall79d07d22020-04-27 18:20:16 +0100217 shrink_axis_mask = op.attrs["shrink_axis_mask"]
Louis Verhaard7db78962020-05-25 15:05:26 +0200218 ellipsis_mask = op.attrs["ellipsis_mask"]
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200219
220 if (new_axis_mask != 0 and shrink_axis_mask != 0) or ellipsis_mask != 0:
221 # Not supported, will be put on CPU
222 return tens
223 if shrink_axis_mask == 0 and new_axis_mask == 0:
Tim Hall79d07d22020-04-27 18:20:16 +0100224 # Equal Rank StridedSlice, no need to insert reshape
225 return tens
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200226 elif shrink_axis_mask != 0:
227 n = 0
228 axis = 0
229 while shrink_axis_mask:
230 prev_mask = shrink_axis_mask
231 n += 1
232 shrink_axis_mask &= shrink_axis_mask - 1
233 axis = int(math.log2(prev_mask - shrink_axis_mask))
234 reshape_input_shape = reshape_input_shape[:axis] + [1] + reshape_input_shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100235
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200236 assert len(tens.shape) == (len(op.inputs[0].shape) - n)
237 op.attrs["shrink_axis_mask"] = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100238
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200239 elif new_axis_mask != 0:
240 n = 0
241 axis = 0
242 while new_axis_mask:
243 prev_mask = new_axis_mask
244 n += 1
245 new_axis_mask &= new_axis_mask - 1
246 axis = int(math.log2(prev_mask - new_axis_mask))
Louis Verhaard7db78962020-05-25 15:05:26 +0200247 reshape_input_shape = reshape_input_shape[:axis] + reshape_input_shape[(axis + 1) :]
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200248 new_axis_mask >>= 1
249
250 assert len(tens.shape) == (len(op.inputs[0].shape) + n)
251 op.attrs["new_axis_mask"] = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100252 else:
253 axis = int(op.attrs["axis"])
254 op.type = "UnpackReshaped"
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200255 reshape_input_shape = tens.shape[:axis] + [1] + tens.shape[axis:]
Tim Hall79d07d22020-04-27 18:20:16 +0100256
257 # Construct 1 shape tensor to be used by all inserted reshape ops
258 new_shape_name = op.name + "_reshape_shape"
259 new_shape_tens = Tensor([1], DataType.int32, new_shape_name)
260 new_shape_tens.values = np.array(tens.shape)
261 new_shape_tens_const = Operation("Const", new_shape_tens.name + "_const")
262 new_shape_tens.ops = [new_shape_tens_const]
263 new_shape_tens_const.outputs = [new_shape_tens]
264
265 for idx, out_tens in enumerate(op.outputs):
266 reshape_name = op.name + str(idx) + "_reshape"
267 reshape_op = Operation("Reshape", reshape_name)
268 reshape_op.outputs = [out_tens]
269 reshape_in = out_tens.clone("_reshaped")
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200270 reshape_in.shape = reshape_in.storage_shape = reshape_in.bandwidth_shape = reshape_input_shape
Tim Hall79d07d22020-04-27 18:20:16 +0100271 reshape_in.ops = [op]
272 out_tens.ops = [reshape_op]
273 reshape_op.inputs = [reshape_in, new_shape_tens]
274
275 op.outputs[idx] = reshape_in
276
277 return tens
278
279
280def add_padding_fields(op, arch):
281 if "padding" in op.attrs:
282 if "Conv" in op.type:
283 kernel_size = op.inputs[1].shape[:2]
284 input_shape = op.inputs[0].shape
Dwight Lidman3ec04ac2020-04-30 11:54:48 +0200285 elif "Pool" in op.type or "ResizeBilinear" == op.type:
Tim Hall79d07d22020-04-27 18:20:16 +0100286 kernel_size = op.attrs["ksize"][1:3]
287 input_shape = op.inputs[0].shape
288 elif op.type == "ExtractImagePatches":
289 kernel_size = op.attrs["ksizes"][1:3]
290 input_shape = op.inputs[0].shape
291 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200292 raise UnsupportedFeatureError("Unknown operation that uses padding: {}".format(op.type))
Tim Hall79d07d22020-04-27 18:20:16 +0100293
294 padding, skirt = calc_padding_and_skirt(op.attrs["padding"], kernel_size, op.attrs["strides"], input_shape)
295 op.attrs["explicit_padding"] = padding
296 op.attrs["skirt"] = skirt
297 return op
298
299
300conv_op = set(("Conv2D", "QuantizedConv2D", "Conv2DBackpropInputSwitched", "Conv2DBiasAct"))
301fc_op = set(
302 (
303 "MatMul",
304 "QuantizedMatMul",
305 "BlockLSTM",
306 "RnnAct",
307 "UnidirectionalSequenceRnnAct",
308 "BidirectionalSequenceRnnAct",
309 "LstmAct",
310 "UnidirectionalSequenceLstmAct",
311 "BidirectionalSequenceLstmAct",
312 "FullyConnectedAct",
313 )
314)
315depthwise_op = set(("DepthwiseConv2dNative", "DepthwiseConv2dBiasAct",))
Louis Verhaard7db78962020-05-25 15:05:26 +0200316pool_op = set(
317 ("AvgPool", "MaxPool", "QuantizedAvgPool", "QuantizedMaxPool", "AvgPoolAct", "MaxPoolAct", "ResizeBilinear",)
318)
Tim Hall79d07d22020-04-27 18:20:16 +0100319elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum", "LeakyRelu", "Abs"))
Charles Xu78792222020-05-13 10:15:26 +0200320binary_elementwise_op = set(("AddAct", "MulAct", "SubAct", "Maximum", "Minimum"))
Tim Hall79d07d22020-04-27 18:20:16 +0100321activation_ops = set(("Relu", "Relu6", "ReluN1To1", "Sigmoid", "Tanh"))
322memory_only_ops = set(("Reshape",))
323
Diego Russoea6111a2020-04-14 18:41:58 +0100324
Tim Hall79d07d22020-04-27 18:20:16 +0100325# Check if the op can be reordered
326def get_prepend_op(op):
327 inp = op.inputs[0]
328 # The op should be reordered between prev_op and prep_op
329 prev_op = inp.ops[-1]
330 prep_op = None
331 while prev_op.type in memory_only_ops and len(prev_op.outputs) == 1 and len(prev_op.outputs[0].consumers()) == 1:
332 prep_op = prev_op
333 inp = prev_op.inputs[0]
334 prev_op = inp.ops[-1]
Diego Russoea6111a2020-04-14 18:41:58 +0100335 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 +0100336 return prep_op
337
338 return None
339
340
341def mark_npu_block_type(op, arch):
342 npu_block_type = NpuBlockType.Default
343 if op.type in conv_op:
344 npu_block_type = NpuBlockType.ConvolutionMxN
345 elif op.type in fc_op:
346 npu_block_type = NpuBlockType.VectorProduct
347 elif op.type in depthwise_op:
348 npu_block_type = NpuBlockType.ConvolutionDepthWise
349 elif op.type in pool_op:
350 npu_block_type = NpuBlockType.Pooling
351 elif op.type in elementwise_op:
352 npu_block_type = NpuBlockType.ElementWise
353
354 op.attrs["npu_block_type"] = npu_block_type
355 return op
356
357
358def convert_depthwise_to_conv(op, arch):
359 # Depthwise is equivalent to a single conv2d if the ifm depth is 1 and
360 # the ofm depth equals the depth multipler.
361 # If those conditions are true, then we can perform a simple
362 # switch of the operator type (and weight order)
363
364 if ("DepthwiseConv2d" in op.type) and (op.attrs["depth_multiplier"] != 1):
365 ifm_tensor = op.inputs[0]
366 weight_tensor = op.inputs[1]
367 ofm_tensor = op.outputs[0]
368 if (ifm_tensor.shape[3] == 1) and (ofm_tensor.shape[3] == op.attrs["depth_multiplier"]):
369 # Change op type to Conv2d
370 op.type = op.type.replace("DepthwiseConv2d", "Conv2D")
371 del op.attrs["channel_multiplier"]
372 del op.attrs["depth_multiplier"]
373
374 weight_tensor.quant_values = np.transpose(weight_tensor.quant_values, (0, 1, 3, 2))
375 weight_tensor.shape = weight_tensor.storage_shape = weight_tensor.bandwidth_shape = list(
376 weight_tensor.quant_values.shape
377 )
378 else:
Louis Verhaard7db78962020-05-25 15:05:26 +0200379 raise UnsupportedFeatureError(
380 "Unsupported DepthwiseConv2d with depth_multiplier = {}, ifm channels = {}, ofm channels = {}".format(
Tim Hall79d07d22020-04-27 18:20:16 +0100381 op.attrs["depth_multiplier"], ifm_tensor.shape[3], ofm_tensor.shape[3]
382 )
383 )
Tim Hall79d07d22020-04-27 18:20:16 +0100384 return op
385
386
387# Reorder activation op if it's after the memory only operations
388def fixup_act_reorder(op, arch):
389 if op.type in activation_ops:
390 prep_op = get_prepend_op(op)
Diego Russoea6111a2020-04-14 18:41:58 +0100391 if prep_op is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100392 act_op = op.clone("_reordered")
393 act_op.inputs = [prep_op.inputs[0]]
394 act_op_out = act_op.inputs[0].clone("_acted")
395 act_op_out.quantization = op.outputs[0].quantization.clone()
396 act_op_out.ops = [act_op]
397 act_op.outputs = [act_op_out]
398 prep_op.inputs[0] = act_op_out
399 prep_op.outputs[0].quantization = act_op_out.quantization.clone()
400
401 # Mark the op so that it will be removed as passthrough later on
402 op.type = "Identity"
403 return op
404
Charles Xu78792222020-05-13 10:15:26 +0200405def fixup_elementwise_with_scalars(op, arch):
406 if op.type in binary_elementwise_op:
407 ifm_tensor, ifm2_tensor, _, ofm_tensor = op.get_ifm_ifm2_weights_ofm()
408 if ifm2_tensor.shape != [] and ifm_tensor.shape != []:
409 diff = len(ifm_tensor.shape) - len(ifm2_tensor.shape)
410 if diff > 0:
411 ifm2_tensor.shape = full_shape(len(ifm_tensor.shape), ifm2_tensor.shape, 1)
412 elif diff < 0:
413 ifm_tensor.shape = full_shape(len(ifm2_tensor.shape), ifm_tensor.shape, 1)
414 return op
Tim Hall79d07d22020-04-27 18:20:16 +0100415
Tim Hall4e127762020-05-15 16:05:49 +0100416# Set input/output tensor equivalence to the same id for memory operations
417def set_tensor_equivalence(op, arch):
418 if op.type == "Reshape":
419 eid = op.outputs[0].equivalence_id
420 for inp in op.inputs:
421 inp.equivalence_id = eid
422 return op
423
424
Tim Hall79d07d22020-04-27 18:20:16 +0100425def convert_mul_max_to_abs_or_lrelu(op, arch):
Diego Russoea6111a2020-04-14 18:41:58 +0100426 r"""Whenever there is a subgraph with this topology:
Tim Hall79d07d22020-04-27 18:20:16 +0100427
428 Input X For X = -1 or X > 0
429 | \ / This subgraph can be replaced with either
430 | Mul an Abs (if X = -1) or a LeakyReLU (if X > 0)
431 | /
432 Max
433 """
434
435 if op.type == "Maximum":
436 # finds the Mul input(s) to the Max
437 muls = [i for i in op.inputs if i.ops[0].type == "MulAct"]
438 if len(muls) == 1:
439 mul = muls[0].ops[0]
440 elif len(muls) == 2:
441 # In the case both inputs are Muls, find the one with the same input as the Max
442 mul = [m for m in muls if len(set(op.inputs + m.ops[0].inputs)) == 1][0].ops[0]
443 else:
444 # No Mul inputs
445 return op
446
447 # make sure the Mul doesn't have any other consumers
448 if len(mul.outputs[0].consumers()) != 1:
449 return op
450 # make sure the Mul doesn't have a faf
451 if mul.attrs["fused_activation_function"]:
452 return op
453
454 # finds the branched input that goes to both the Max and the Mul
455 shared = set(op.inputs) & set(mul.inputs)
456 if len(shared) == 1:
457 shared_in = shared.pop()
458 # find the constant scalar input to the Mul
459 const_tens = (set(mul.inputs) - {shared_in}).pop()
460 # check that it is a scalar
461 if const_tens.shape != []:
462 return op
463 const = const_tens.ops[0]
464 # check that it is a constant
465 if const.type != "Const":
466 return op
467 else:
468 return op
469
470 val = const.outputs[0].values
471 if val >= 0:
472 new_op = "LeakyRelu"
473 op.attrs["alpha"] = val
474 elif val == -1:
475 new_op = "Abs"
476 else:
477 return op
478
479 op.type = op.type.replace("Maximum", new_op)
480 op.name = op.name.replace("Maximum", new_op)
481 op.outputs[0].name = op.outputs[0].name.replace("Maximum", new_op)
482 op.inputs = [shared_in]
483 return op
484
485
486def supported_operator_check(op, arch):
487 op.run_on_npu = arch.supported_operators.is_operator_supported(op)
488 return op
489
490
491def optimise_graph_a(nng, arch, verbose_graph=False):
492 if verbose_graph:
493 nng.print_graph()
494
495 op_rewrite_list = [
496 # mark block type and check if the operations are supported
497 mark_npu_block_type,
Tim Hall4e127762020-05-15 16:05:49 +0100498 set_tensor_equivalence,
Tim Hall79d07d22020-04-27 18:20:16 +0100499 supported_operator_check,
500 # then do any rewrites of supported operators
501 convert_depthwise_to_conv,
502 fixup_fully_connected_input,
503 fixup_pack_input,
504 fixup_conv2d_backprop,
505 fixup_act_reorder,
506 add_padding_fields,
507 mark_npu_block_type,
Charles Xu78792222020-05-13 10:15:26 +0200508 fixup_elementwise_with_scalars,
Tim Hall79d07d22020-04-27 18:20:16 +0100509 # convert_mul_max_to_abs_or_lrelu # TODO: enable optimisation once quantisation issues are resolved
510 ]
511
512 for idx, sg in enumerate(nng.subgraphs):
513 # rewrite graph pass
514 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(
Diego Russoea6111a2020-04-14 18:41:58 +0100515 sg, arch, [fixup_unpack_output], op_rewrite_list, rewrite_unsupported=False
Tim Hall79d07d22020-04-27 18:20:16 +0100516 )
517
518 for idx, sg in enumerate(nng.subgraphs):
519 # remove passthrough tensors
Diego Russoea6111a2020-04-14 18:41:58 +0100520 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(sg, arch, [remove_passthrough_tensor], [])
Tim Hall79d07d22020-04-27 18:20:16 +0100521
522 if verbose_graph:
523 nng.print_graph()
524 return nng
525
Diego Russoea6111a2020-04-14 18:41:58 +0100526
Tim Hall79d07d22020-04-27 18:20:16 +0100527def optimise_graph_b(nng, arch, verbose_graph=False):
528 if verbose_graph:
529 nng.print_graph()
530
531 for idx, sg in enumerate(nng.subgraphs):
532 # combined rewrite graph pass
Diego Russoea6111a2020-04-14 18:41:58 +0100533 nng.subgraphs[idx] = rewrite_graph.rewrite_graph_pre_order(sg, arch, [rewrite_concat, rewrite_split], [])
Tim Hall79d07d22020-04-27 18:20:16 +0100534
535 if verbose_graph:
536 nng.print_graph()
537 return nng