blob: 705f839b46ade7cc70613e7e6a1de522d17a4ccb [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# Mark purpose and select formats for Tensors. Also compresses the weights.
Tim Hall79d07d22020-04-27 18:20:16 +010018from . import rewrite_graph
19from . import weight_compressor
Tim Hallc8310b12020-06-17 14:53:11 +010020from .errors import OperatorError
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020021from .tensor import MemType
Diego Russoe8a10452020-04-21 17:39:10 +010022from .tensor import TensorFormat
23from .tensor import TensorPurpose
Tim Hallc8310b12020-06-17 14:53:11 +010024from .tflite_mapping import custom_prefix
Tim Hall79d07d22020-04-27 18:20:16 +010025
26
27def purpose_from_list(lst):
28 def purpose(op, idx):
29 return lst[idx]
30
31 return purpose
32
33
34def all_fm(op, idx):
35 return TensorPurpose.FeatureMap
36
37
38def all_parameter(op, idx):
39 return TensorPurpose.FeatureMap
40
41
42def input0_from_output_rest_parameter(op, idx):
43 if idx == 0:
44 res = op.outputs[0].purpose
45 if res == TensorPurpose.Unknown:
46 print("Warning: Propagating unknown tensor purpose", op)
47 return res
48 return TensorPurpose.FeatureMap
49
50
51def inputs_from_output(op, idx):
52 res = op.outputs[0].purpose
53 if res == TensorPurpose.Unknown:
54 print("Warning: Propagating unknown tensor purpose", op)
55 return res
56
Diego Russoea6111a2020-04-14 18:41:58 +010057
Tim Hall79d07d22020-04-27 18:20:16 +010058tensor_purposes = [ # ops, input_purpose
59 (
60 set(
61 (
62 "Relu",
63 "Relu6",
64 "Mul",
65 "Add",
66 "Sub",
67 "Rsqrt",
68 "Abs",
69 "Cast",
70 "Exp",
71 "Floor",
72 "FloorDiv",
73 "FloorMod",
74 "SquaredDifference",
75 "AddN",
76 "BiasAdd",
77 "RealDiv",
78 "Maximum",
79 "Minimum",
80 "Sigmoid",
81 "Tanh",
82 "FusedBatchNorm",
83 "AvgPool",
84 "MaxPool",
85 "Squeeze",
86 "Softmax",
87 "LRN",
88 "Assign",
89 "BatchMatMul",
90 "ZerosLike",
91 "ExtractImagePatches",
92 "MulAct",
93 "AddAct",
94 "SubAct",
95 "DivAct",
96 "AvgPoolAct",
97 "MaxPoolAct",
98 "LeakyRelu",
99 )
100 ),
101 all_fm,
102 ),
103 (
104 set(
105 (
106 "Conv2D",
107 "DepthwiseConv2dNative",
108 "MatMul",
109 "Conv2DBiasAct",
110 "DepthwiseConv2dBiasAct",
111 "FullyConnectedAct",
112 )
113 ),
114 purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap]),
115 ),
116 (
Jacob Bohlincf7da102020-05-20 09:03:40 +0200117 set(("Conv2DBackpropInputSwitchedBias",)),
Tim Hallc30f4952020-06-15 20:47:35 +0100118 purpose_from_list(
119 [TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]
120 ),
Tim Hall79d07d22020-04-27 18:20:16 +0100121 ),
122 (
123 set(("QuantizedConv2D", "QuantizedMatMul")),
124 purpose_from_list(
125 [
126 TensorPurpose.FeatureMap,
127 TensorPurpose.Weights,
128 TensorPurpose.FeatureMap,
129 TensorPurpose.FeatureMap,
130 TensorPurpose.FeatureMap,
131 TensorPurpose.FeatureMap,
132 ]
133 ),
134 ),
135 (
136 set(
137 (
138 "Reshape",
139 "Min",
140 "Max",
141 "Mean",
142 "Pad",
143 "MirrorPad",
144 "ArgMax",
145 "ArgMin",
146 "ExpandDims",
147 "ResizeNearestNeighbor",
148 "ResizeBilinear",
149 "Tile",
150 "Transpose",
151 "Mfcc",
152 )
153 ),
154 purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]),
155 ),
156 (
157 set(("QuantizedReshape", "QuantizedResizeBilinear")),
158 purpose_from_list(
159 [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]
160 ),
161 ),
162 (
163 set(("QuantizedBiasAdd", "QuantizedAdd", "QuantizedMul")),
164 purpose_from_list(
165 [
166 TensorPurpose.FeatureMap,
167 TensorPurpose.FeatureMap,
168 TensorPurpose.FeatureMap,
169 TensorPurpose.FeatureMap,
170 TensorPurpose.FeatureMap,
171 TensorPurpose.FeatureMap,
172 ]
173 ),
174 ),
175 (
176 set(
177 (
178 "Dequantize",
179 "Quantize",
180 "QuantizeV2",
181 "QuantizedRelu",
182 "QuantizedRelu1",
183 "QuantizedRelu6",
184 "QuantizedAvgPool",
185 "QuantizedMaxPool",
186 "Slice",
187 "SplitV",
188 )
189 ),
190 purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]),
191 ),
192 (
193 set(("BatchToSpaceND", "SpaceToBatchND", "DepthToSpaceND", "SpaceToDepthND")),
194 purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]),
195 ),
196 (
197 set(("BlockLSTM",)),
198 purpose_from_list(
199 [
200 TensorPurpose.FeatureMap,
201 TensorPurpose.FeatureMap,
202 TensorPurpose.FeatureMap,
203 TensorPurpose.FeatureMap,
204 TensorPurpose.Weights,
205 TensorPurpose.FeatureMap,
206 TensorPurpose.FeatureMap,
207 TensorPurpose.FeatureMap,
208 TensorPurpose.FeatureMap,
209 ]
210 ),
211 ),
212 (set(("SplitSliceRead",)), purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap])),
213 (set(("Shape", "ConcatSliceWrite", "AudioSpectrogram")), purpose_from_list([TensorPurpose.FeatureMap])),
214 (
215 set(("StridedSlice",)),
216 purpose_from_list(
217 [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]
218 ),
219 ),
220 (set(("Fill", "Pack", "Range")), all_parameter),
221 (
222 set(("Requantize",)),
223 purpose_from_list(
224 [
225 TensorPurpose.FeatureMap,
226 TensorPurpose.FeatureMap,
227 TensorPurpose.FeatureMap,
228 TensorPurpose.FeatureMap,
229 TensorPurpose.FeatureMap,
230 ]
231 ),
232 ),
233 (set(("Placeholder", "SubgraphInput", "Const", "VariableV2")), purpose_from_list([])),
234 (set(("FakeQuantWithMinMaxArgs", "FakeQuantWithMinMaxVars")), input0_from_output_rest_parameter),
235 (
236 set(("Square", "Sqrt", "Log", "Less", "Enter", "Exit", "Identity", "StopGradient", "Merge", "Switch")),
237 inputs_from_output,
238 ),
239 (None, all_fm),
240]
241
242
243for ops, input_purpose in tensor_purposes:
244 if ops is None:
245 continue
246 for op in ops:
247 assert len(op) > 1, "string literal has been decomposed"
248
249
250def mark_tensor_purpose(nng, arch, verbose_tensor_purpose=False):
251 def mark_tensor_helper(tens, purpose):
252
253 if tens.purpose == TensorPurpose.Unknown or tens.purpose == purpose:
254 tens.purpose = purpose
255 else:
256 assert 0, "Cannot resolve tensor purpose %s and %s for tensor %s" % (tens.purpose, purpose, tens)
257 tens.mem_area = arch.tensor_storage_mem_area[tens.purpose]
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200258 tens.mem_type = arch.tensor_storage_mem_type[tens.purpose]
Tim Hall79d07d22020-04-27 18:20:16 +0100259
260 if len(tens.ops) == 1 and tens.ops[0].type == "Const":
261 tens.mem_area = (
262 arch.permanent_storage_mem_area
263 ) # special case constants, as they must be in permanent storage
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200264 tens.mem_type = MemType.Permanent_NPU
Tim Hall79d07d22020-04-27 18:20:16 +0100265
266 def rewrite_mark_tensor_purpose(op, arch):
267 # find disconnected outputs and mark as parameters
268 for tens in op.outputs:
269 if not tens.consumers():
270 mark_tensor_helper(tens, TensorPurpose.FeatureMap)
271
272 for ops, input_purpose in tensor_purposes:
273 if ops is None or op.type in ops:
274 if ops is None:
275 print(
Tim Hallc8310b12020-06-17 14:53:11 +0100276 "Warning: Don't know how to mark up purpose for",
Tim Hall79d07d22020-04-27 18:20:16 +0100277 op.type,
278 op.inputs,
279 "triggering all feature map fallback",
280 )
Tim Hallc8310b12020-06-17 14:53:11 +0100281
Tim Hall79d07d22020-04-27 18:20:16 +0100282 for idx, tens in enumerate(op.inputs):
283 purpose = input_purpose(op, idx)
284 mark_tensor_helper(tens, purpose)
Tim Hallc8310b12020-06-17 14:53:11 +0100285
Louis Verhaardc4cbbc92020-05-18 13:40:02 +0200286 if op.type == "Reshape":
287 # Reshape's input and output point to same data
288 op.outputs[0].mem_area = op.inputs[0].mem_area
Tim Hallc8310b12020-06-17 14:53:11 +0100289
290 if op.type.startswith(custom_prefix) and op.attrs.get("custom_type", "") == "ExistingNpuOp":
291 scratch_tensor = None
292
293 if len(op.inputs) >= 3:
294 scratch_tensor = op.inputs[2] # should be existing scratch tensor
295 if scratch_tensor.name.endswith("_scratch"):
296 scratch_tensor.purpose = TensorPurpose.Scratch
297
298 if scratch_tensor is None:
299 raise OperatorError(op, "Scratch tensor not found.")
300
Tim Hall79d07d22020-04-27 18:20:16 +0100301 break
Tim Hallc8310b12020-06-17 14:53:11 +0100302
Tim Hall79d07d22020-04-27 18:20:16 +0100303 return op
304
305 for sg in nng.subgraphs:
306 sg = rewrite_graph.rewrite_graph_pre_order(sg, arch, [], [rewrite_mark_tensor_purpose])
307 for tens in sg.output_tensors:
308 mark_tensor_helper(tens, TensorPurpose.FeatureMap)
309
310 if verbose_tensor_purpose:
311 nng.print_graph_with_tensors()
312
313 return nng
314
315
316reshape_operations = set(
317 (
318 "Reshape",
319 "QuantizedReshape",
320 "ExpandDims",
321 "Squeeze",
322 "BatchToSpaceND",
323 "SpaceToBatchND",
324 "DepthToSpaceND",
325 "SpaceToDepthND",
326 "Placeholder",
327 )
328)
329
330
331def mark_tensor_format(nng, arch, verbose_tensor_format=False):
332 formats_for_tensor = {}
333
334 def init_tens(tens):
335 if tens.purpose == TensorPurpose.FeatureMap:
336 fmt = arch.default_feature_map_format
337 elif tens.purpose == TensorPurpose.Weights:
338 fmt = arch.default_weight_format
Tim Hallc8310b12020-06-17 14:53:11 +0100339 elif tens.purpose == TensorPurpose.Scratch:
340 fmt = arch.default_feature_map_format
Tim Hall465582c2020-05-26 09:33:14 +0100341 elif tens.purpose == TensorPurpose.Unknown:
342 fmt = TensorFormat.Unknown
Tim Hall79d07d22020-04-27 18:20:16 +0100343 else:
344 assert 0, "unknown tensor purpose %s" % (tens.purpose,)
345 return fmt
346
Tim Hall79d07d22020-04-27 18:20:16 +0100347 def visit_tens(tens, ps):
Diego Russoea6111a2020-04-14 18:41:58 +0100348 if tens not in formats_for_tensor:
Tim Hall79d07d22020-04-27 18:20:16 +0100349 fmt = init_tens(tens)
350 else:
351 fmt = formats_for_tensor[tens]
352
353 formats_for_tensor[tens] = fmt
354
355 for sg in nng.subgraphs:
356 for ps in sg.passes:
357 for tens in ps.outputs:
358 visit_tens(tens, ps)
359 for tens in ps.intermediates:
360 visit_tens(tens, ps)
361 for tens in ps.inputs:
362 visit_tens(tens, ps)
363
364 for tens, fmt in formats_for_tensor.items():
365 tens.set_format(fmt, arch)
366 if fmt == TensorFormat.WeightsCompressed and tens.values is not None:
Louis Verhaard3c07c972020-05-07 08:12:58 +0200367 src_tens = tens.get_dma_src_tensor()
368 if src_tens is not None:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200369 op = tens.find_npu_op()
370 npu_block_type = op.attrs["npu_block_type"]
371 weight_compressor.compress_weights(arch, nng, tens, npu_block_type, 32, 32, op.get_dilation_h_w())
Tim Hall79d07d22020-04-27 18:20:16 +0100372 # Alias compressed weights back into source tensor
Louis Verhaard3c07c972020-05-07 08:12:58 +0200373 src_tens.copy_compressed_weight_info(tens)
Tim Hall79d07d22020-04-27 18:20:16 +0100374
375 if verbose_tensor_format:
376 nng.print_passes_with_tensors()