Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 1 | # 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 Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 16 | # Description: |
| 17 | # Mark purpose and select formats for Tensors. Also compresses the weights. |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 18 | from . import rewrite_graph |
| 19 | from . import weight_compressor |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 20 | from .errors import OperatorError |
Patrik Gustavsson | eca2e95 | 2020-05-27 09:15:11 +0200 | [diff] [blame] | 21 | from .tensor import MemType |
Diego Russo | e8a1045 | 2020-04-21 17:39:10 +0100 | [diff] [blame] | 22 | from .tensor import TensorFormat |
| 23 | from .tensor import TensorPurpose |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 24 | from .tflite_mapping import custom_prefix |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 25 | |
| 26 | |
| 27 | def purpose_from_list(lst): |
| 28 | def purpose(op, idx): |
| 29 | return lst[idx] |
| 30 | |
| 31 | return purpose |
| 32 | |
| 33 | |
| 34 | def all_fm(op, idx): |
| 35 | return TensorPurpose.FeatureMap |
| 36 | |
| 37 | |
| 38 | def all_parameter(op, idx): |
| 39 | return TensorPurpose.FeatureMap |
| 40 | |
| 41 | |
| 42 | def 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 | |
| 51 | def 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 Russo | ea6111a | 2020-04-14 18:41:58 +0100 | [diff] [blame] | 57 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 58 | tensor_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", |
Fredrik Svedberg | a0c3624 | 2020-06-03 15:43:31 +0200 | [diff] [blame] | 99 | "CLZ", |
| 100 | "SHL", |
| 101 | "SHR", |
| 102 | "ReduceSum", |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 103 | ) |
| 104 | ), |
| 105 | all_fm, |
| 106 | ), |
| 107 | ( |
| 108 | set( |
| 109 | ( |
| 110 | "Conv2D", |
| 111 | "DepthwiseConv2dNative", |
| 112 | "MatMul", |
| 113 | "Conv2DBiasAct", |
| 114 | "DepthwiseConv2dBiasAct", |
| 115 | "FullyConnectedAct", |
| 116 | ) |
| 117 | ), |
| 118 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap]), |
| 119 | ), |
| 120 | ( |
Jacob Bohlin | cf7da10 | 2020-05-20 09:03:40 +0200 | [diff] [blame] | 121 | set(("Conv2DBackpropInputSwitchedBias",)), |
Tim Hall | c30f495 | 2020-06-15 20:47:35 +0100 | [diff] [blame] | 122 | purpose_from_list( |
| 123 | [TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap] |
| 124 | ), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 125 | ), |
| 126 | ( |
| 127 | set(("QuantizedConv2D", "QuantizedMatMul")), |
| 128 | purpose_from_list( |
| 129 | [ |
| 130 | TensorPurpose.FeatureMap, |
| 131 | TensorPurpose.Weights, |
| 132 | TensorPurpose.FeatureMap, |
| 133 | TensorPurpose.FeatureMap, |
| 134 | TensorPurpose.FeatureMap, |
| 135 | TensorPurpose.FeatureMap, |
| 136 | ] |
| 137 | ), |
| 138 | ), |
| 139 | ( |
| 140 | set( |
| 141 | ( |
| 142 | "Reshape", |
| 143 | "Min", |
| 144 | "Max", |
| 145 | "Mean", |
| 146 | "Pad", |
| 147 | "MirrorPad", |
| 148 | "ArgMax", |
| 149 | "ArgMin", |
| 150 | "ExpandDims", |
| 151 | "ResizeNearestNeighbor", |
| 152 | "ResizeBilinear", |
| 153 | "Tile", |
| 154 | "Transpose", |
| 155 | "Mfcc", |
| 156 | ) |
| 157 | ), |
| 158 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]), |
| 159 | ), |
| 160 | ( |
| 161 | set(("QuantizedReshape", "QuantizedResizeBilinear")), |
| 162 | purpose_from_list( |
| 163 | [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap] |
| 164 | ), |
| 165 | ), |
| 166 | ( |
| 167 | set(("QuantizedBiasAdd", "QuantizedAdd", "QuantizedMul")), |
| 168 | purpose_from_list( |
| 169 | [ |
| 170 | TensorPurpose.FeatureMap, |
| 171 | TensorPurpose.FeatureMap, |
| 172 | TensorPurpose.FeatureMap, |
| 173 | TensorPurpose.FeatureMap, |
| 174 | TensorPurpose.FeatureMap, |
| 175 | TensorPurpose.FeatureMap, |
| 176 | ] |
| 177 | ), |
| 178 | ), |
| 179 | ( |
| 180 | set( |
| 181 | ( |
| 182 | "Dequantize", |
| 183 | "Quantize", |
| 184 | "QuantizeV2", |
| 185 | "QuantizedRelu", |
| 186 | "QuantizedRelu1", |
| 187 | "QuantizedRelu6", |
| 188 | "QuantizedAvgPool", |
| 189 | "QuantizedMaxPool", |
| 190 | "Slice", |
| 191 | "SplitV", |
| 192 | ) |
| 193 | ), |
| 194 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]), |
| 195 | ), |
| 196 | ( |
| 197 | set(("BatchToSpaceND", "SpaceToBatchND", "DepthToSpaceND", "SpaceToDepthND")), |
| 198 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]), |
| 199 | ), |
| 200 | ( |
| 201 | set(("BlockLSTM",)), |
| 202 | purpose_from_list( |
| 203 | [ |
| 204 | TensorPurpose.FeatureMap, |
| 205 | TensorPurpose.FeatureMap, |
| 206 | TensorPurpose.FeatureMap, |
| 207 | TensorPurpose.FeatureMap, |
| 208 | TensorPurpose.Weights, |
| 209 | TensorPurpose.FeatureMap, |
| 210 | TensorPurpose.FeatureMap, |
| 211 | TensorPurpose.FeatureMap, |
| 212 | TensorPurpose.FeatureMap, |
| 213 | ] |
| 214 | ), |
| 215 | ), |
| 216 | (set(("SplitSliceRead",)), purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap])), |
| 217 | (set(("Shape", "ConcatSliceWrite", "AudioSpectrogram")), purpose_from_list([TensorPurpose.FeatureMap])), |
| 218 | ( |
| 219 | set(("StridedSlice",)), |
| 220 | purpose_from_list( |
| 221 | [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap] |
| 222 | ), |
| 223 | ), |
| 224 | (set(("Fill", "Pack", "Range")), all_parameter), |
| 225 | ( |
| 226 | set(("Requantize",)), |
| 227 | purpose_from_list( |
| 228 | [ |
| 229 | TensorPurpose.FeatureMap, |
| 230 | TensorPurpose.FeatureMap, |
| 231 | TensorPurpose.FeatureMap, |
| 232 | TensorPurpose.FeatureMap, |
| 233 | TensorPurpose.FeatureMap, |
| 234 | ] |
| 235 | ), |
| 236 | ), |
| 237 | (set(("Placeholder", "SubgraphInput", "Const", "VariableV2")), purpose_from_list([])), |
| 238 | (set(("FakeQuantWithMinMaxArgs", "FakeQuantWithMinMaxVars")), input0_from_output_rest_parameter), |
| 239 | ( |
| 240 | set(("Square", "Sqrt", "Log", "Less", "Enter", "Exit", "Identity", "StopGradient", "Merge", "Switch")), |
| 241 | inputs_from_output, |
| 242 | ), |
| 243 | (None, all_fm), |
| 244 | ] |
| 245 | |
| 246 | |
| 247 | for ops, input_purpose in tensor_purposes: |
| 248 | if ops is None: |
| 249 | continue |
| 250 | for op in ops: |
| 251 | assert len(op) > 1, "string literal has been decomposed" |
| 252 | |
| 253 | |
| 254 | def mark_tensor_purpose(nng, arch, verbose_tensor_purpose=False): |
| 255 | def mark_tensor_helper(tens, purpose): |
| 256 | |
| 257 | if tens.purpose == TensorPurpose.Unknown or tens.purpose == purpose: |
| 258 | tens.purpose = purpose |
Fredrik Svedberg | a0c3624 | 2020-06-03 15:43:31 +0200 | [diff] [blame] | 259 | elif tens.purpose != TensorPurpose.LUT: |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 260 | assert 0, "Cannot resolve tensor purpose %s and %s for tensor %s" % (tens.purpose, purpose, tens) |
| 261 | tens.mem_area = arch.tensor_storage_mem_area[tens.purpose] |
Patrik Gustavsson | eca2e95 | 2020-05-27 09:15:11 +0200 | [diff] [blame] | 262 | tens.mem_type = arch.tensor_storage_mem_type[tens.purpose] |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 263 | |
| 264 | if len(tens.ops) == 1 and tens.ops[0].type == "Const": |
| 265 | tens.mem_area = ( |
| 266 | arch.permanent_storage_mem_area |
| 267 | ) # special case constants, as they must be in permanent storage |
Patrik Gustavsson | eca2e95 | 2020-05-27 09:15:11 +0200 | [diff] [blame] | 268 | tens.mem_type = MemType.Permanent_NPU |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 269 | |
| 270 | def rewrite_mark_tensor_purpose(op, arch): |
| 271 | # find disconnected outputs and mark as parameters |
| 272 | for tens in op.outputs: |
| 273 | if not tens.consumers(): |
| 274 | mark_tensor_helper(tens, TensorPurpose.FeatureMap) |
| 275 | |
| 276 | for ops, input_purpose in tensor_purposes: |
| 277 | if ops is None or op.type in ops: |
| 278 | if ops is None: |
| 279 | print( |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 280 | "Warning: Don't know how to mark up purpose for", |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 281 | op.type, |
| 282 | op.inputs, |
| 283 | "triggering all feature map fallback", |
| 284 | ) |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 285 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 286 | for idx, tens in enumerate(op.inputs): |
| 287 | purpose = input_purpose(op, idx) |
| 288 | mark_tensor_helper(tens, purpose) |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 289 | |
Louis Verhaard | c4cbbc9 | 2020-05-18 13:40:02 +0200 | [diff] [blame] | 290 | if op.type == "Reshape": |
| 291 | # Reshape's input and output point to same data |
| 292 | op.outputs[0].mem_area = op.inputs[0].mem_area |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 293 | |
| 294 | if op.type.startswith(custom_prefix) and op.attrs.get("custom_type", "") == "ExistingNpuOp": |
| 295 | scratch_tensor = None |
| 296 | |
| 297 | if len(op.inputs) >= 3: |
| 298 | scratch_tensor = op.inputs[2] # should be existing scratch tensor |
| 299 | if scratch_tensor.name.endswith("_scratch"): |
| 300 | scratch_tensor.purpose = TensorPurpose.Scratch |
| 301 | |
| 302 | if scratch_tensor is None: |
| 303 | raise OperatorError(op, "Scratch tensor not found.") |
| 304 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 305 | break |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 306 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 307 | return op |
| 308 | |
| 309 | for sg in nng.subgraphs: |
| 310 | sg = rewrite_graph.rewrite_graph_pre_order(sg, arch, [], [rewrite_mark_tensor_purpose]) |
| 311 | for tens in sg.output_tensors: |
| 312 | mark_tensor_helper(tens, TensorPurpose.FeatureMap) |
| 313 | |
| 314 | if verbose_tensor_purpose: |
| 315 | nng.print_graph_with_tensors() |
| 316 | |
| 317 | return nng |
| 318 | |
| 319 | |
| 320 | reshape_operations = set( |
| 321 | ( |
| 322 | "Reshape", |
| 323 | "QuantizedReshape", |
| 324 | "ExpandDims", |
| 325 | "Squeeze", |
| 326 | "BatchToSpaceND", |
| 327 | "SpaceToBatchND", |
| 328 | "DepthToSpaceND", |
| 329 | "SpaceToDepthND", |
| 330 | "Placeholder", |
| 331 | ) |
| 332 | ) |
| 333 | |
| 334 | |
| 335 | def mark_tensor_format(nng, arch, verbose_tensor_format=False): |
| 336 | formats_for_tensor = {} |
| 337 | |
| 338 | def init_tens(tens): |
Fredrik Svedberg | a0c3624 | 2020-06-03 15:43:31 +0200 | [diff] [blame] | 339 | if tens.purpose in (TensorPurpose.FeatureMap, TensorPurpose.LUT): |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 340 | fmt = arch.default_feature_map_format |
| 341 | elif tens.purpose == TensorPurpose.Weights: |
| 342 | fmt = arch.default_weight_format |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 343 | elif tens.purpose == TensorPurpose.Scratch: |
| 344 | fmt = arch.default_feature_map_format |
Tim Hall | 465582c | 2020-05-26 09:33:14 +0100 | [diff] [blame] | 345 | elif tens.purpose == TensorPurpose.Unknown: |
| 346 | fmt = TensorFormat.Unknown |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 347 | else: |
| 348 | assert 0, "unknown tensor purpose %s" % (tens.purpose,) |
| 349 | return fmt |
| 350 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 351 | def visit_tens(tens, ps): |
Diego Russo | ea6111a | 2020-04-14 18:41:58 +0100 | [diff] [blame] | 352 | if tens not in formats_for_tensor: |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 353 | fmt = init_tens(tens) |
| 354 | else: |
| 355 | fmt = formats_for_tensor[tens] |
| 356 | |
| 357 | formats_for_tensor[tens] = fmt |
| 358 | |
| 359 | for sg in nng.subgraphs: |
| 360 | for ps in sg.passes: |
| 361 | for tens in ps.outputs: |
| 362 | visit_tens(tens, ps) |
| 363 | for tens in ps.intermediates: |
| 364 | visit_tens(tens, ps) |
| 365 | for tens in ps.inputs: |
| 366 | visit_tens(tens, ps) |
| 367 | |
| 368 | for tens, fmt in formats_for_tensor.items(): |
| 369 | tens.set_format(fmt, arch) |
| 370 | if fmt == TensorFormat.WeightsCompressed and tens.values is not None: |
Louis Verhaard | 3c07c97 | 2020-05-07 08:12:58 +0200 | [diff] [blame] | 371 | src_tens = tens.get_dma_src_tensor() |
| 372 | if src_tens is not None: |
Louis Verhaard | b2fb212 | 2020-06-04 15:51:24 +0200 | [diff] [blame] | 373 | op = tens.find_npu_op() |
Dwight Lidman | 940fdee | 2020-08-13 13:11:48 +0200 | [diff] [blame] | 374 | if op is not None: |
| 375 | npu_block_type = op.attrs["npu_block_type"] |
| 376 | weight_compressor.compress_weights(arch, nng, tens, npu_block_type, 16, 16, op.get_dilation_h_w()) |
| 377 | # Alias compressed weights back into source tensor |
| 378 | src_tens.copy_compressed_weight_info(tens) |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 379 | |
| 380 | if verbose_tensor_format: |
| 381 | nng.print_passes_with_tensors() |