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 |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 21 | from .operation import CustomType |
| 22 | from .operation import Op |
Patrik Gustavsson | eca2e95 | 2020-05-27 09:15:11 +0200 | [diff] [blame] | 23 | from .tensor import MemType |
Diego Russo | e8a1045 | 2020-04-21 17:39:10 +0100 | [diff] [blame] | 24 | from .tensor import TensorFormat |
| 25 | from .tensor import TensorPurpose |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 26 | |
| 27 | |
| 28 | def purpose_from_list(lst): |
| 29 | def purpose(op, idx): |
| 30 | return lst[idx] |
| 31 | |
| 32 | return purpose |
| 33 | |
| 34 | |
| 35 | def all_fm(op, idx): |
| 36 | return TensorPurpose.FeatureMap |
| 37 | |
| 38 | |
| 39 | def all_parameter(op, idx): |
| 40 | return TensorPurpose.FeatureMap |
| 41 | |
| 42 | |
| 43 | def input0_from_output_rest_parameter(op, idx): |
| 44 | if idx == 0: |
| 45 | res = op.outputs[0].purpose |
| 46 | if res == TensorPurpose.Unknown: |
| 47 | print("Warning: Propagating unknown tensor purpose", op) |
| 48 | return res |
| 49 | return TensorPurpose.FeatureMap |
| 50 | |
| 51 | |
| 52 | def inputs_from_output(op, idx): |
| 53 | res = op.outputs[0].purpose |
| 54 | if res == TensorPurpose.Unknown: |
| 55 | print("Warning: Propagating unknown tensor purpose", op) |
| 56 | return res |
| 57 | |
Diego Russo | ea6111a | 2020-04-14 18:41:58 +0100 | [diff] [blame] | 58 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 59 | tensor_purposes = [ # ops, input_purpose |
| 60 | ( |
| 61 | set( |
| 62 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 63 | Op.Relu, |
| 64 | Op.Relu6, |
| 65 | Op.Rsqrt, |
| 66 | Op.Abs, |
| 67 | Op.Cast, |
| 68 | Op.Exp, |
| 69 | Op.Floor, |
| 70 | Op.FloorDiv, |
| 71 | Op.FloorMod, |
| 72 | Op.SquaredDifference, |
| 73 | Op.AddN, |
| 74 | Op.Maximum, |
| 75 | Op.Minimum, |
| 76 | Op.Sigmoid, |
| 77 | Op.Tanh, |
| 78 | Op.AvgPool, |
| 79 | Op.MaxPool, |
| 80 | Op.Squeeze, |
| 81 | Op.Softmax, |
| 82 | Op.LRN, |
| 83 | Op.BatchMatMul, |
| 84 | Op.ZerosLike, |
| 85 | Op.Mul, |
| 86 | Op.Add, |
| 87 | Op.Sub, |
| 88 | Op.Div, |
| 89 | Op.LeakyRelu, |
| 90 | Op.CLZ, |
| 91 | Op.SHL, |
| 92 | Op.SHR, |
| 93 | Op.ReduceSum, |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 94 | ) |
| 95 | ), |
| 96 | all_fm, |
| 97 | ), |
| 98 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 99 | set((Op.Conv2D, Op.MatMul, Op.Conv2DBias, Op.DepthwiseConv2DBias, Op.FullyConnected,)), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 100 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap]), |
| 101 | ), |
| 102 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 103 | set((Op.Conv2DBackpropInputSwitchedBias,)), |
Tim Hall | c30f495 | 2020-06-15 20:47:35 +0100 | [diff] [blame] | 104 | purpose_from_list( |
| 105 | [TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap] |
| 106 | ), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 107 | ), |
| 108 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 109 | set((Op.QuantizedConv2D, Op.QuantizedMatMul)), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 110 | purpose_from_list( |
| 111 | [ |
| 112 | TensorPurpose.FeatureMap, |
| 113 | TensorPurpose.Weights, |
| 114 | TensorPurpose.FeatureMap, |
| 115 | TensorPurpose.FeatureMap, |
| 116 | TensorPurpose.FeatureMap, |
| 117 | TensorPurpose.FeatureMap, |
| 118 | ] |
| 119 | ), |
| 120 | ), |
| 121 | ( |
| 122 | set( |
| 123 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 124 | Op.Reshape, |
| 125 | Op.Min, |
| 126 | Op.Max, |
| 127 | Op.Mean, |
| 128 | Op.Pad, |
| 129 | Op.MirrorPad, |
| 130 | Op.ArgMax, |
| 131 | Op.ArgMin, |
| 132 | Op.ExpandDims, |
| 133 | Op.ResizeNearestNeighbor, |
| 134 | Op.ResizeBilinear, |
| 135 | Op.Tile, |
| 136 | Op.Transpose, |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 137 | ) |
| 138 | ), |
| 139 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]), |
| 140 | ), |
| 141 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 142 | set((Op.QuantizedReshape,)), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 143 | purpose_from_list( |
| 144 | [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap] |
| 145 | ), |
| 146 | ), |
| 147 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 148 | set((Op.Dequantize, Op.Quantize, Op.QuantizedAvgPool, Op.QuantizedMaxPool, Op.Slice, Op.SplitV,)), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 149 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]), |
| 150 | ), |
| 151 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 152 | set((Op.BatchToSpaceND, Op.SpaceToBatchND, Op.DepthToSpace, Op.SpaceToDepth)), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 153 | purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]), |
| 154 | ), |
| 155 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 156 | set((Op.BlockLSTM,)), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 157 | purpose_from_list( |
| 158 | [ |
| 159 | TensorPurpose.FeatureMap, |
| 160 | TensorPurpose.FeatureMap, |
| 161 | TensorPurpose.FeatureMap, |
| 162 | TensorPurpose.FeatureMap, |
| 163 | TensorPurpose.Weights, |
| 164 | TensorPurpose.FeatureMap, |
| 165 | TensorPurpose.FeatureMap, |
| 166 | TensorPurpose.FeatureMap, |
| 167 | TensorPurpose.FeatureMap, |
| 168 | ] |
| 169 | ), |
| 170 | ), |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 171 | (set((Op.SplitSliceRead,)), purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap])), |
| 172 | (set((Op.Shape, Op.ConcatSliceWrite)), purpose_from_list([TensorPurpose.FeatureMap])), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 173 | ( |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 174 | set((Op.StridedSlice,)), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 175 | purpose_from_list( |
| 176 | [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap] |
| 177 | ), |
| 178 | ), |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 179 | (set((Op.Fill, Op.Pack, Op.Range)), all_parameter), |
| 180 | (set((Op.Placeholder, Op.SubgraphInput, Op.Const,)), purpose_from_list([])), |
| 181 | (set((Op.FakeQuantWithMinMaxArgs,)), input0_from_output_rest_parameter), |
| 182 | (set((Op.Square, Op.Sqrt, Op.Log, Op.Less, Op.Identity,)), inputs_from_output,), |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 183 | (None, all_fm), |
| 184 | ] |
| 185 | |
| 186 | |
| 187 | for ops, input_purpose in tensor_purposes: |
| 188 | if ops is None: |
| 189 | continue |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 190 | |
| 191 | |
| 192 | def mark_tensor_purpose(nng, arch, verbose_tensor_purpose=False): |
| 193 | def mark_tensor_helper(tens, purpose): |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 194 | if tens.purpose == TensorPurpose.Unknown or tens.purpose == purpose: |
| 195 | tens.purpose = purpose |
Fredrik Svedberg | a0c3624 | 2020-06-03 15:43:31 +0200 | [diff] [blame] | 196 | elif tens.purpose != TensorPurpose.LUT: |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 197 | assert 0, "Cannot resolve tensor purpose %s and %s for tensor %s" % (tens.purpose, purpose, tens) |
| 198 | tens.mem_area = arch.tensor_storage_mem_area[tens.purpose] |
Patrik Gustavsson | eca2e95 | 2020-05-27 09:15:11 +0200 | [diff] [blame] | 199 | tens.mem_type = arch.tensor_storage_mem_type[tens.purpose] |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 200 | |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 201 | if len(tens.ops) == 1 and tens.ops[0].type == Op.Const: |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 202 | tens.mem_area = ( |
| 203 | arch.permanent_storage_mem_area |
| 204 | ) # special case constants, as they must be in permanent storage |
Patrik Gustavsson | eca2e95 | 2020-05-27 09:15:11 +0200 | [diff] [blame] | 205 | tens.mem_type = MemType.Permanent_NPU |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 206 | |
Patrik Gustavsson | 3010d9b | 2020-10-01 08:22:10 +0200 | [diff] [blame] | 207 | def rewrite_mark_tensor_purpose(op, arch, nng): |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 208 | # find disconnected outputs and mark as parameters |
| 209 | for tens in op.outputs: |
| 210 | if not tens.consumers(): |
| 211 | mark_tensor_helper(tens, TensorPurpose.FeatureMap) |
| 212 | |
| 213 | for ops, input_purpose in tensor_purposes: |
| 214 | if ops is None or op.type in ops: |
| 215 | if ops is None: |
| 216 | print( |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 217 | "Warning: Don't know how to mark up purpose for", |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 218 | op.type, |
| 219 | op.inputs, |
| 220 | "triggering all feature map fallback", |
| 221 | ) |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 222 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 223 | for idx, tens in enumerate(op.inputs): |
Andreas Nevalainen | d8c032d | 2020-09-11 10:25:09 +0200 | [diff] [blame] | 224 | if tens is None: |
| 225 | continue |
Louis Verhaard | b9fc33c | 2020-08-13 11:47:36 +0200 | [diff] [blame] | 226 | purpose = input_purpose(op, idx) if tens.purpose == TensorPurpose.Unknown else tens.purpose |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 227 | mark_tensor_helper(tens, purpose) |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 228 | |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 229 | if op.type == Op.Reshape: |
Louis Verhaard | c4cbbc9 | 2020-05-18 13:40:02 +0200 | [diff] [blame] | 230 | # Reshape's input and output point to same data |
| 231 | op.outputs[0].mem_area = op.inputs[0].mem_area |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 232 | |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 233 | if op.type == Op.Custom and op.attrs.get("custom_type") == CustomType.ExistingNpuOp: |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 234 | scratch_tensor = None |
| 235 | |
| 236 | if len(op.inputs) >= 3: |
| 237 | scratch_tensor = op.inputs[2] # should be existing scratch tensor |
| 238 | if scratch_tensor.name.endswith("_scratch"): |
| 239 | scratch_tensor.purpose = TensorPurpose.Scratch |
| 240 | |
| 241 | if scratch_tensor is None: |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 242 | OperatorError(op, "Scratch tensor not found.") |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 243 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 244 | break |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 245 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 246 | return op |
| 247 | |
| 248 | for sg in nng.subgraphs: |
Patrik Gustavsson | 3010d9b | 2020-10-01 08:22:10 +0200 | [diff] [blame] | 249 | sg = rewrite_graph.rewrite_graph_pre_order(nng, sg, arch, [], [rewrite_mark_tensor_purpose]) |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 250 | for tens in sg.output_tensors: |
| 251 | mark_tensor_helper(tens, TensorPurpose.FeatureMap) |
| 252 | |
| 253 | if verbose_tensor_purpose: |
| 254 | nng.print_graph_with_tensors() |
| 255 | |
| 256 | return nng |
| 257 | |
| 258 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 259 | def mark_tensor_format(nng, arch, verbose_tensor_format=False): |
| 260 | formats_for_tensor = {} |
| 261 | |
| 262 | def init_tens(tens): |
Fredrik Svedberg | a0c3624 | 2020-06-03 15:43:31 +0200 | [diff] [blame] | 263 | if tens.purpose in (TensorPurpose.FeatureMap, TensorPurpose.LUT): |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 264 | fmt = arch.default_feature_map_format |
| 265 | elif tens.purpose == TensorPurpose.Weights: |
| 266 | fmt = arch.default_weight_format |
Tim Hall | c8310b1 | 2020-06-17 14:53:11 +0100 | [diff] [blame] | 267 | elif tens.purpose == TensorPurpose.Scratch: |
| 268 | fmt = arch.default_feature_map_format |
Tim Hall | 465582c | 2020-05-26 09:33:14 +0100 | [diff] [blame] | 269 | elif tens.purpose == TensorPurpose.Unknown: |
| 270 | fmt = TensorFormat.Unknown |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 271 | else: |
| 272 | assert 0, "unknown tensor purpose %s" % (tens.purpose,) |
| 273 | return fmt |
| 274 | |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 275 | def visit_tens(tens, ps): |
Diego Russo | ea6111a | 2020-04-14 18:41:58 +0100 | [diff] [blame] | 276 | if tens not in formats_for_tensor: |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 277 | fmt = init_tens(tens) |
| 278 | else: |
| 279 | fmt = formats_for_tensor[tens] |
| 280 | |
| 281 | formats_for_tensor[tens] = fmt |
| 282 | |
| 283 | for sg in nng.subgraphs: |
| 284 | for ps in sg.passes: |
| 285 | for tens in ps.outputs: |
| 286 | visit_tens(tens, ps) |
| 287 | for tens in ps.intermediates: |
| 288 | visit_tens(tens, ps) |
| 289 | for tens in ps.inputs: |
| 290 | visit_tens(tens, ps) |
| 291 | |
| 292 | for tens, fmt in formats_for_tensor.items(): |
Fredrik Svedberg | 0f98b36 | 2020-09-29 10:00:39 +0200 | [diff] [blame] | 293 | if len(tens.shape) > 4: |
| 294 | continue |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 295 | tens.set_format(fmt, arch) |
| 296 | if fmt == TensorFormat.WeightsCompressed and tens.values is not None: |
Louis Verhaard | 3c07c97 | 2020-05-07 08:12:58 +0200 | [diff] [blame] | 297 | src_tens = tens.get_dma_src_tensor() |
| 298 | if src_tens is not None: |
Louis Verhaard | b2fb212 | 2020-06-04 15:51:24 +0200 | [diff] [blame] | 299 | op = tens.find_npu_op() |
Dwight Lidman | 940fdee | 2020-08-13 13:11:48 +0200 | [diff] [blame] | 300 | if op is not None: |
Louis Verhaard | aee5d75 | 2020-09-30 09:01:52 +0200 | [diff] [blame^] | 301 | weight_compressor.compress_weights( |
| 302 | arch, nng, tens, op.type.npu_block_type, 16, 16, op.get_dilation_h_w() |
| 303 | ) |
Dwight Lidman | 940fdee | 2020-08-13 13:11:48 +0200 | [diff] [blame] | 304 | # Alias compressed weights back into source tensor |
| 305 | src_tens.copy_compressed_weight_info(tens) |
Tim Hall | 79d07d2 | 2020-04-27 18:20:16 +0100 | [diff] [blame] | 306 | |
| 307 | if verbose_tensor_format: |
| 308 | nng.print_passes_with_tensors() |