blob: b07b4dc3fd4c8ac15eda449cff229940881c76c6 [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# Internal representation of a Neural Network Tensor.
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +010018import copy
Tim Hall79d07d22020-04-27 18:20:16 +010019import enum
Tim Hall79d07d22020-04-27 18:20:16 +010020import uuid
Jacob Bohlin1a666972020-09-11 10:04:15 +020021from collections import defaultdict
Louis Verhaard9db529a2020-09-23 10:27:11 +020022from functools import lru_cache
Diego Russoea6111a2020-04-14 18:41:58 +010023
24import numpy as np
25
26from . import numeric_util
Tim Hall93582962020-09-09 21:58:15 +010027from .data_type import BaseType
Michael McGeagh5778ffd2020-08-06 17:31:02 +010028from .data_type import DataType
Dwight Lidmana9390f72020-05-13 12:00:08 +020029from .ethos_u55_regs.ethos_u55_regs import resampling_mode
Louis Verhaardaee5d752020-09-30 09:01:52 +020030from .operation import Op
Michael McGeagh5778ffd2020-08-06 17:31:02 +010031from .operation import Operation
Diego Russoe8a10452020-04-21 17:39:10 +010032from .range_set import MemoryRangeSet
Tim Hall79d07d22020-04-27 18:20:16 +010033
34
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020035class MemType(enum.IntFlag):
36 Unknown = 0
37 Permanent_NPU = 1
38 Permanent_CPU = 2
39 Scratch = 3
40 Scratch_fast = 4
41 Size = Scratch_fast + 1
42
43 def display_name(self):
44 return ("Unknown", "Permanent_NPU", "Permanent_CPU", "Scratch", "Scratch_fast", "Size")[self.value]
45
46 def identifier_name(self):
47 return ("unknown", "permanent_npu", "permanent_cpu", "scratch", "scratch_fast", "size")[self.value]
48
49 def all():
50 return (MemType.Permanent_NPU, MemType.Permanent_CPU, MemType.Scratch, MemType.Scratch_fast)
51
52 def __str__(self):
53 return self.name
54
55
Tim Hall79d07d22020-04-27 18:20:16 +010056class MemArea(enum.IntFlag):
57 Unknown = 0
58 Sram = 1
59 Dram = 2
60 OnChipFlash = 3
61 OffChipFlash = 4
Louis Verhaard0b8268a2020-08-05 16:11:29 +020062 Shram = 5 # for LUT
63 Size = Shram + 1
Tim Hall79d07d22020-04-27 18:20:16 +010064
65 def display_name(self):
Louis Verhaard0b8268a2020-08-05 16:11:29 +020066 return ("Unknown", "SRAM", "DRAM", "On-chip Flash", "Off-chip Flash", "SHRAM", "Size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +010067
68 def identifier_name(self):
Louis Verhaard0b8268a2020-08-05 16:11:29 +020069 return ("unknown", "sram", "dram", "on_chip_flash", "off_chip_flash", "shram", "size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +010070
71 def all():
Louis Verhaard0b8268a2020-08-05 16:11:29 +020072 return (MemArea.Sram, MemArea.Dram, MemArea.OnChipFlash, MemArea.OffChipFlash, MemArea.Shram)
Tim Hall79d07d22020-04-27 18:20:16 +010073
74 def __str__(self):
75 return self.name
76
77
78class TensorPurpose(enum.IntFlag):
79 Unknown = 0
80 Weights = 1
81 FeatureMap = 2
82 Scratch = 3
Fredrik Svedberga0c36242020-06-03 15:43:31 +020083 LUT = 4
Andreas Nevalainen897cc142020-10-28 15:42:08 +010084 FSBias = 5
85 Size = 6
Tim Hall79d07d22020-04-27 18:20:16 +010086
87 def display_name(self):
Andreas Nevalainen897cc142020-10-28 15:42:08 +010088 return ("Unknown", "Weights", "FeatureMap", "Scratch", "LUT", "FastStorageBias", "Size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +010089
90 def identifier_name(self):
Andreas Nevalainen897cc142020-10-28 15:42:08 +010091 return ("unknown", "weights", "feature_map", "scratch", "lut", "fast_storage_bias", "size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +010092
93 def all():
Andreas Nevalainen897cc142020-10-28 15:42:08 +010094 return (TensorPurpose.Weights, TensorPurpose.FeatureMap, TensorPurpose.FSBias)
Tim Hall79d07d22020-04-27 18:20:16 +010095
96
97class TensorSubPurpose(enum.Enum):
98 Standard = 0
99 DoubleBuffer = 1
100 RollingBufferX = 2
101 RollingBufferY = 3
102 RollingBufferXY = 4
103
104 def display_name(self):
105 return ("Standard", "Double Buffer", "Rolling Buffer X", "Rolling Buffer Y", "Rolling Buffer XY")[self.value]
106
107 def identifier_name(self):
108 return ("standard", "double_buffer", "rolling_buffer_x", "rolling_buffer_y", "rolling_buffer_xy")[self.value]
109
110 def all():
111 return (
112 TensorSubPurpose.Standard,
113 TensorSubPurpose.DoubleBuffer,
114 TensorSubPurpose.RollingBufferX,
115 TensorSubPurpose.RollingBufferY,
116 TensorSubPurpose.RollingBufferXY,
117 )
118
119
120class TensorFormat(enum.Flag):
121 Unknown = 0
122 WeightsCompressed = 1
123 NHWC = 2
124 NHCWB16 = 3
125
126 def __str__(self):
127 return self.name
128
129
130class TensorBlockTraversal(enum.Enum):
131 Default = 0
132 DepthWise = 1
133 DepthFirst = 2
134 PartKernelFirst = 3
135
136
137def shape_num_elements(shp):
138 elems = 1
139 if shp is None:
140 return None
141 for d in shp:
142 if d is None:
143 return None
144 elems *= d
145 return elems
146
147
148def shape_fully_defined(shp):
149 if shp is None:
150 return False
151 for d in shp:
152 if d is None:
153 return False
154 return True
155
156
157def shape_round_to_quantum(shp, quantum):
158 new_shp = list(shp)
159
160 # Traverse backwards using length of shape since there may be more rounding quantums than shape elements
161 for i in range(-1, -len(shp) - 1, -1):
162 if new_shp[i] is not None:
163 new_shp[i] = numeric_util.round_up(new_shp[i], quantum[i])
164 return new_shp
165
166
Louis Verhaard9db529a2020-09-23 10:27:11 +0200167@lru_cache(maxsize=None)
168def create_equivalence_id(key):
169 # Generates equivalence_id based on the given key.
170 return uuid.uuid4()
171
172
Tim Hall79d07d22020-04-27 18:20:16 +0100173class QuantizationParameters:
174 __slots__ = "min", "max", "num_bits", "narrow_range", "scale_f32", "zero_point", "quant_min", "quant_max"
175
176 def __init__(self, min=None, max=None, num_bits=None, narrow_range=None):
177 self.min = min
178 self.max = max
179
180 self.num_bits = num_bits
181 self.narrow_range = narrow_range
182
183 self.scale_f32 = None
184 self.zero_point = None
185 self.quant_min = None
186 self.quant_max = None
187
188 def __str__(self):
189 return "<nng.QuantizationParameters min=%s max=%s, num_bits=%s, scale=%s, zero_point=%s>" % (
190 self.min,
191 self.max,
192 self.num_bits,
193 self.scale_f32,
194 self.zero_point,
195 )
196
197 __repr__ = __str__
198
199 def clone(self):
200 res = QuantizationParameters()
201 res.min = self.min
202 res.max = self.max
203
204 res.num_bits = self.num_bits
205 res.narrow_range = self.narrow_range
206
207 res.scale_f32 = self.scale_f32
208 res.zero_point = self.zero_point
209 res.quant_min = self.quant_min
210 res.quant_max = self.quant_max
211 return res
212
213 def dequantize(self, values):
214 if self.zero_point.size == 1 and self.scale_f32.size == 1:
215 # same scale is used for all values
216 res = (values.astype(np.float64) - self.zero_point) * self.scale_f32
217 else:
218 # a different scale is used for different sets of values
219 values_as_float = values.astype(np.float64)
220
221 # this is not compatible with the format of depthwise weights,
222 # where input is at index 3 (Output, Kh, Kw, Input)
223 # return the quantized values
224 return np.ndarray((values_as_float.shape))
225
226 shape = values_as_float.shape[0]
227 assert self.zero_point.size == self.scale_f32.size == shape
228 res = np.ndarray(values_as_float.shape)
229 for i in range(shape):
230 res[i] = (values_as_float[i] - self.zero_point[i]) * self.scale_f32[i]
231
232 return res
233
Tim Halle3786ac2020-07-28 17:40:50 +0100234 def is_scaling_equal(self, other):
Tim Hall93582962020-09-09 21:58:15 +0100235 # quantisation parameter scaling is not equal if 'other' is None because
236 # it implies that the tensor it belongs to is not quantised. otherwise,
237 # it depends upon whether the scale and zero point are equal
238
Tim Hall89567612020-10-27 11:57:57 +0000239 if not isinstance(other, QuantizationParameters):
Tim Halle3786ac2020-07-28 17:40:50 +0100240 return False
241
242 return self.scale_f32 == other.scale_f32 and self.zero_point == other.zero_point
243
Tim Hall93582962020-09-09 21:58:15 +0100244 def is_valid(self):
245 # quantisation parameters are consider valid if they have a scale and zero point
246
247 return None not in (self.scale_f32, self.zero_point)
248
Dwight Lidmanc7187432020-11-16 17:40:46 +0100249 def is_per_axis(self):
250 """Returns True if either the scale, zero point, minimum or maximum values are arrays"""
251 for attr in ("scale_f32", "zero_point", "min", "max"):
252 if isinstance(getattr(self, attr), np.ndarray):
253 return True
254 return False
255
Tim Hall79d07d22020-04-27 18:20:16 +0100256
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100257def create_const_tensor(name, shape, dtype, values, value_dtype=None, purpose=TensorPurpose.Unknown, quantization=None):
258 # Tensor
259 const_tensor = Tensor(shape, dtype, name + "_0")
260 const_tensor.purpose = purpose
261 const_tensor.quantization = quantization
262 const_tensor.values = np.array(values, dtype=value_dtype)
Jacob Bohlina41cd4d2020-08-26 18:21:28 +0200263 const_tensor.quant_values = np.frombuffer(const_tensor.values.tobytes(), dtype=np.uint8)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100264 # Operator
Louis Verhaardaee5d752020-09-30 09:01:52 +0200265 const_op = Operation(Op.Const, name)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100266 const_op.set_output_tensor(const_tensor)
267 return const_tensor
268
269
270def create_reshape_tensor(tens, shape, ifm_reshape=True):
271 if shape == tens.shape:
272 return tens
273 # Tensors
274 name = tens.name + "_reshape"
275 reshape_ifm = tens
276 reshape_ofm = tens.clone("_reshaped")
277 reshape_ofm.set_all_shapes(shape)
278 if not ifm_reshape:
279 reshape_ifm, reshape_ofm = reshape_ofm, reshape_ifm
280 # Operator
Louis Verhaardaee5d752020-09-30 09:01:52 +0200281 reshape_op = Operation(Op.Reshape, name)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100282 reshape_op.attrs["new_shape"] = shape
283 reshape_op.add_input_tensor(reshape_ifm)
284 reshape_op.add_input_tensor(create_const_tensor(name + "_shape", [1], DataType.int32, shape))
285 reshape_op.set_output_tensor(reshape_ofm)
286 return reshape_ofm if ifm_reshape else reshape_ifm
287
288
Jacob Bohlin1a666972020-09-11 10:04:15 +0200289# class that keeps track of all tensor addresses in the different memory types
290class TensorAddressMap:
291 address_map = defaultdict(dict) # dict (tens.equivalence_id -> dict (mem_type -> address))
292
293 @classmethod
294 def get_address_for_tens(cls, tens_id, mem_type):
295 return cls.address_map[tens_id].get(mem_type)
296
297 @classmethod
298 def set_address_for_tens(cls, tens_id, mem_type, address):
299 # Check previous address if there is one
300 previous_address = cls.address_map[tens_id].get(mem_type)
Louis Verhaard0b9c9a32020-09-15 14:05:38 +0200301 if address is not None and previous_address is not None:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200302 assert previous_address == address, "Two different addresses cannot be assigned to the same tensor."
303
304 # Set tensor's address for memory type
305 cls.address_map[tens_id][mem_type] = address
306
307
Tim Hall79d07d22020-04-27 18:20:16 +0100308class Tensor:
309 __slots__ = (
310 "shape",
311 "storage_shape",
312 "bandwidth_shape",
313 "dtype",
314 "name",
315 "ops",
316 "consumer_list",
317 "values",
318 "quant_values",
319 "compressed_values",
Tim Hallf7e810a2020-06-25 15:04:31 +0100320 "compressed_values_substream_offsets",
Tim Hall79d07d22020-04-27 18:20:16 +0100321 "mem_area",
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200322 "mem_type",
Tim Hall79d07d22020-04-27 18:20:16 +0100323 "format",
324 "purpose",
325 "sub_purpose",
326 "alignment",
327 "weight_transpose_depthwise",
328 "storage_compression_scale",
329 "bandwidth_compression_scale",
330 "compression_scale_for_worst_weight_stream",
331 "weight_compression_scales",
332 "weight_compression_config",
Louis Verhaard9db529a2020-09-23 10:27:11 +0200333 "value_id",
Tim Hall79d07d22020-04-27 18:20:16 +0100334 "storage_rounding_quantum",
335 "brick_size",
Tim Hall79d07d22020-04-27 18:20:16 +0100336 "quantization",
337 "weight_compressed_offsets",
338 "element_size_bytes",
Tim Hall79d07d22020-04-27 18:20:16 +0100339 "block_traversal",
Tim Hall79d07d22020-04-27 18:20:16 +0100340 "equivalence_id",
Dwight Lidmana9390f72020-05-13 12:00:08 +0200341 "resampling_mode",
Patrik Gustavsson458a2082020-08-13 13:41:05 +0200342 "avoid_NHCWB16",
Tim Hall79d07d22020-04-27 18:20:16 +0100343 )
344 AllocationQuantum = 16
345
346 def __init__(self, shape, dtype, name):
347 self.shape = shape
348 self.storage_shape = shape
349 self.bandwidth_shape = shape
350 self.dtype = dtype
351 self.name = name
352 self.equivalence_id = uuid.uuid4()
353
354 self.ops = []
355 self.consumer_list = []
Tim Hall79d07d22020-04-27 18:20:16 +0100356
357 self.values = None
358 self.quant_values = None
359 self.compressed_values = None
Tim Hallf7e810a2020-06-25 15:04:31 +0100360 self.compressed_values_substream_offsets = None
Tim Hall79d07d22020-04-27 18:20:16 +0100361 self.mem_area = MemArea.Unknown
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200362 self.mem_type = MemType.Unknown
Tim Hall79d07d22020-04-27 18:20:16 +0100363 self.format = TensorFormat.Unknown
364 self.purpose = TensorPurpose.Unknown
365 self.sub_purpose = TensorSubPurpose.Standard
366 self.alignment = Tensor.AllocationQuantum
367 self.weight_transpose_depthwise = False
368
369 self.storage_compression_scale = 1.0
370 self.bandwidth_compression_scale = 1.0
371 self.compression_scale_for_worst_weight_stream = 1.0
372 self.weight_compression_scales = None
Louis Verhaard9db529a2020-09-23 10:27:11 +0200373 # if two tensors have the same weight_compression_config, then they have the same compressed values
Tim Hall79d07d22020-04-27 18:20:16 +0100374 self.weight_compression_config = None
Louis Verhaard9db529a2020-09-23 10:27:11 +0200375 # if two tensors have the same value_id, then they have the same values
376 self.value_id = uuid.uuid4()
Tim Hall79d07d22020-04-27 18:20:16 +0100377 self.weight_compressed_offsets = []
378 self.storage_rounding_quantum = (1, 1, 1, 1)
379 self.brick_size = (1, 1, 1, 1)
Tim Hall79d07d22020-04-27 18:20:16 +0100380 self.element_size_bytes = 0
381
382 # quantization parameters
383 self.quantization = None
Tim Hall79d07d22020-04-27 18:20:16 +0100384 self.block_traversal = TensorBlockTraversal.Default
Dwight Lidmana9390f72020-05-13 12:00:08 +0200385 self.resampling_mode = resampling_mode.NONE
Tim Hall79d07d22020-04-27 18:20:16 +0100386
Patrik Gustavsson458a2082020-08-13 13:41:05 +0200387 self.avoid_NHCWB16 = False
388
Jacob Bohlin1a666972020-09-11 10:04:15 +0200389 @property
390 def address(self):
391 return TensorAddressMap.get_address_for_tens(self.equivalence_id, self.mem_type)
392
393 @address.setter
394 def address(self, address):
395 TensorAddressMap.set_address_for_tens(self.equivalence_id, self.mem_type, address)
396
Tim Hall79d07d22020-04-27 18:20:16 +0100397 def element_size(self):
398 if self.element_size_bytes == 0:
399 return self.dtype.size_in_bits() / 8
400 return self.element_size_bytes
401
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100402 # Returns a copy, renamed to self.name + suffix
403 # The references to Operators will be empty when returned
404 # Depending on set_unique, the copy is shallow, or deep
405 # For set_unique==True, a new equivalence_id will be set
406 def clone(self, suffix="_clone", set_unique=False):
407 if set_unique:
408 res = copy.deepcopy(self)
409 res.equivalence_id = uuid.uuid4()
410 else:
411 res = copy.copy(self)
412 res.storage_shape = list(self.storage_shape)
413 res.bandwidth_shape = list(self.bandwidth_shape)
414 if self.quantization is not None:
415 res.quantization = self.quantization.clone()
Tim Hall79d07d22020-04-27 18:20:16 +0100416
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100417 res.name = res.name + suffix
Tim Hall79d07d22020-04-27 18:20:16 +0100418 res.ops = []
419 res.consumer_list = []
Tim Hall79d07d22020-04-27 18:20:16 +0100420
Tim Hall79d07d22020-04-27 18:20:16 +0100421 return res
422
423 def clone_into_fast_storage(self, arch):
424 res = self.clone(suffix="_fast_storage")
425 res.mem_area = arch.fast_storage_mem_area
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200426 res.mem_type = MemType.Scratch_fast
Tim Hall79d07d22020-04-27 18:20:16 +0100427 return res
428
Louis Verhaard3c07c972020-05-07 08:12:58 +0200429 def copy_compressed_weight_info(self, src_tens):
430 # Copies compressed values + all related weight compression info from the given tensor
Louis Verhaard9db529a2020-09-23 10:27:11 +0200431 self.equivalence_id = src_tens.equivalence_id
Louis Verhaard3c07c972020-05-07 08:12:58 +0200432 self.compressed_values = src_tens.compressed_values
Tim Hallf7e810a2020-06-25 15:04:31 +0100433 self.compressed_values_substream_offsets = src_tens.compressed_values_substream_offsets
Louis Verhaard3c07c972020-05-07 08:12:58 +0200434 self.storage_shape = src_tens.storage_shape
435 self.brick_size = src_tens.brick_size
436 self.weight_compression_scales = src_tens.weight_compression_scales
437 self.weight_compressed_offsets = src_tens.weight_compressed_offsets
438 self.weight_transpose_depthwise = src_tens.weight_transpose_depthwise
439 self.compression_scale_for_worst_weight_stream = src_tens.compression_scale_for_worst_weight_stream
440 self.storage_compression_scale = src_tens.storage_compression_scale
Diqing Zhong7e1d1d12020-10-30 15:10:46 +0100441 self.bandwidth_compression_scale = src_tens.bandwidth_compression_scale
Louis Verhaard3c07c972020-05-07 08:12:58 +0200442 self.block_traversal = src_tens.block_traversal
443 self.weight_compression_config = src_tens.weight_compression_config
Louis Verhaard9db529a2020-09-23 10:27:11 +0200444 self.value_id = src_tens.value_id
Louis Verhaard3c07c972020-05-07 08:12:58 +0200445
Tim Hall79d07d22020-04-27 18:20:16 +0100446 def set_format(self, fmt, arch):
447 self.format = fmt
448 shape_len = 0
449 try:
450 shape_len = len(self.shape)
451 except TypeError:
452 pass
453
Louis Verhaard0411edb2020-11-16 16:37:11 +0100454 if shape_len > 4:
455 return
Tim Hall79d07d22020-04-27 18:20:16 +0100456 self.storage_rounding_quantum = arch.storage_rounding_quantums[self.format]
457 self.storage_rounding_quantum = self.storage_rounding_quantum[-shape_len:]
Tim Hall79d07d22020-04-27 18:20:16 +0100458 self.brick_size = arch.brick_sizes[self.format]
459 self.brick_size = self.brick_size[-shape_len:]
460 if self.shape is None:
461 return
462
463 self.bandwidth_shape = shape_round_to_quantum(self.shape, self.brick_size)
464 self.storage_shape = shape_round_to_quantum(self.shape, self.storage_rounding_quantum)
465
466 if fmt == TensorFormat.WeightsCompressed:
467 compression_ratio = 5 / 8
468 self.storage_compression_scale = compression_ratio
469 self.bandwidth_compression_scale = compression_ratio
470 self.compression_scale_for_worst_weight_stream = compression_ratio
471
472 def storage_elements(self):
473 elems = shape_num_elements(self.storage_shape)
474 if elems is None:
475 return 0
476 return elems
477
478 def elements(self):
479 elems = shape_num_elements(self.shape)
480 if elems is None:
481 return 0
482 return elems
483
484 def has_fully_defined_shape(self):
485 return shape_fully_defined(self.shape)
486
Patrik Gustavsson90831bc2020-08-24 16:26:11 +0200487 def storage_size(self, scale=1.0):
488 raw_size = self.storage_elements() * self.element_size() * scale
Tim Hall79d07d22020-04-27 18:20:16 +0100489 if raw_size == 0:
490 raw_size = 1 # force it to take up space
491 rounded_size = numeric_util.round_up(numeric_util.round_up_to_int(raw_size), self.alignment)
492 return rounded_size
493
Patrik Gustavsson90831bc2020-08-24 16:26:11 +0200494 def storage_size_for_sub_purpose(self, arch, sub_purpose, param_a=None, param_b=None):
Tim Hall79d07d22020-04-27 18:20:16 +0100495 alt_shape = self.storage_shape_for_sub_purpose(sub_purpose, param_a, param_b)
496 elems = shape_num_elements(alt_shape)
497 if elems is None:
498 return 0
499 if sub_purpose == TensorSubPurpose.DoubleBuffer:
Patrik Gustavsson90831bc2020-08-24 16:26:11 +0200500 raw_size = (
501 elems
502 * self.element_size()
503 * self.compression_scale_for_worst_weight_stream
504 * arch.weight_estimation_scaling
505 )
Tim Hall79d07d22020-04-27 18:20:16 +0100506 else:
Patrik Gustavsson9baa4c32020-08-20 13:59:01 +0200507 # Rolling buffers are used for intermediate data in ifm streaming
508 # These will all use the NHCWB16 format, and need to be aligned to 16 in the C-dimension
509 if alt_shape[-1] % 16 != 0:
510 nhcwb16_shape = alt_shape[0:-1] + [numeric_util.round_up(alt_shape[-1], 16)]
511 elems = shape_num_elements(nhcwb16_shape)
512
Tim Hall79d07d22020-04-27 18:20:16 +0100513 raw_size = elems * self.element_size() * self.storage_compression_scale
514 rounded_size = numeric_util.round_up(numeric_util.round_up_to_int(raw_size), self.alignment)
515 return rounded_size
516
517 def storage_shape_for_sub_purpose(self, sub_purpose, param_a, param_b):
Tim Hall79d07d22020-04-27 18:20:16 +0100518 if sub_purpose == TensorSubPurpose.DoubleBuffer:
Jacob Bohline843d332020-06-23 12:12:56 +0200519 shp = list(self.shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100520 assert len(shp) >= 2
521 shp[-1] = min(shp[-1], param_a * 2)
Tim Hall79d07d22020-04-27 18:20:16 +0100522 else:
Jacob Bohline843d332020-06-23 12:12:56 +0200523 shp = list(self.storage_shape)
524 if sub_purpose == TensorSubPurpose.RollingBufferX:
525 assert len(shp) == 4
526 shp[0] = 1
527 shp[2] = min(shp[2], param_a)
528 elif sub_purpose == TensorSubPurpose.RollingBufferY:
529 assert len(shp) == 4
530 shp[0] = 1
531 shp[1] = min(shp[1], param_a)
532 elif sub_purpose == TensorSubPurpose.RollingBufferXY:
533 assert len(shp) == 4
534 shp[0] = 1
535 shp[2] = min(shp[2], param_a)
536 shp[1] = min(shp[1], param_b)
537 elif sub_purpose == TensorSubPurpose.Standard:
538 pass
539 else:
540 assert 0, "did not expect new sub purpose %s" % (sub_purpose,)
541
Tim Hall79d07d22020-04-27 18:20:16 +0100542 return shp
543
544 def set_new_sub_purpose(self, sub_purpose, param_a=None, param_b=None):
545 self.storage_shape = self.storage_shape_for_sub_purpose(sub_purpose, param_a, param_b)
546 self.sub_purpose = sub_purpose
547 if sub_purpose == TensorSubPurpose.DoubleBuffer:
548 self.storage_compression_scale = self.compression_scale_for_worst_weight_stream
549
550 def bandwidth(self):
551 elems = shape_num_elements(self.bandwidth_shape)
552 if elems is None:
553 return 0
554 return elems * self.element_size() * self.bandwidth_compression_scale
555
556 def consumers(self):
557 return self.consumer_list
558
559 def get_address_ranges_for_coordinates(self, start_coord, end_coord):
560 if self.sub_purpose in set(
561 (TensorSubPurpose.RollingBufferX, TensorSubPurpose.RollingBufferY, TensorSubPurpose.RollingBufferXY)
562 ):
563 # build dummy coordinates that cover the entire buffer
564 start_coord = [0] * len(start_coord)
565 end_coord = [min(self.storage_shape[i], self.shape[i]) for i in range(len(end_coord))]
566
567 start = self.address_for_coordinate(start_coord, is_top_box=False)
568 end = self.address_for_coordinate(end_coord, is_top_box=True)
569 return MemoryRangeSet(self.mem_area, start, end)
570
571 def addresses_for_rolling_buffer(self, start_coord, end_coord):
572 # returns ( box_height0, box_height1, box_width, [address_tl, address_tr, address_bl, address_br] )
573
574 if len(start_coord) < 4:
575 box_height0 = 1
576 box_width = 1
577
578 if len(start_coord) >= 2:
579 box_width = end_coord[-2] - start_coord[-2]
580
581 return box_height0, box_height0, box_width, [self.address_for_coordinate(start_coord), None, None, None]
582
583 crossing_y = numeric_util.round_up(start_coord[1] + 1, self.storage_shape[1])
584 crossing_x = numeric_util.round_up(start_coord[2] + 1, self.storage_shape[2])
585
586 crossing_y = min(crossing_y, end_coord[1])
587 crossing_x = min(crossing_x, end_coord[2])
588
589 box_height0 = crossing_y - start_coord[1]
590 box_width = crossing_x - start_coord[2]
591
592 addresses = [None] * 4
593 addresses[0] = self.address_for_coordinate(start_coord)
594
595 if end_coord[2] > crossing_x:
596 addresses[1] = self.address_for_coordinate([start_coord[0], start_coord[1], crossing_x, start_coord[3]])
597 raise Exception("Striping in vertical direction is not supported")
598 if end_coord[1] > crossing_y:
599 addresses[2] = self.address_for_coordinate([start_coord[0], crossing_y, start_coord[2], start_coord[3]])
600 if end_coord[1] > crossing_y and end_coord[2] > crossing_x:
601 addresses[3] = self.address_for_coordinate([start_coord[0], crossing_y, crossing_x, start_coord[3]])
602
603 return box_height0, box_height0, box_width, addresses
604
605 def address_for_coordinate(self, coord, is_top_box=False):
606 return self.address + self.address_offset_for_coordinate(coord, is_top_box)
607
608 def get_strides_and_coord(self, coord=None):
609 if coord is None:
610 coord = [0] * len(self.storage_shape)
611
612 augmented_coord = coord
613 augmented_shape = self.storage_shape
614 while len(augmented_shape) < 4:
615 augmented_shape = [1] + augmented_shape
616
617 while len(augmented_coord) < 4:
618 augmented_coord = [0] + augmented_coord
619
620 assert len(augmented_coord) == len(augmented_shape)
621
622 if self.format == TensorFormat.NHWC:
623 augmented_shape = [augmented_shape[0], augmented_shape[3]] + augmented_shape[1:3] + [1]
624 augmented_coord = [augmented_coord[0], augmented_coord[3]] + augmented_coord[1:3] + [0]
625 stride_order = [4, 1, 3, 2, 0]
626
627 elif self.format == TensorFormat.NHCWB16:
Patrik Gustavsson2213e902020-05-05 17:49:35 +0200628 channel_divisor = 16
Tim Hall79d07d22020-04-27 18:20:16 +0100629 augmented_shape = augmented_shape[0:4] + [1]
630 augmented_coord = (
631 [augmented_coord[0], augmented_coord[3] // channel_divisor]
632 + augmented_coord[1:3]
633 + [augmented_coord[3] % channel_divisor]
634 )
635
636 if augmented_shape[1] == 0:
637 augmented_shape[1] = 1
638
639 else:
640 assert self.format in set((TensorFormat.Unknown, TensorFormat.WeightsCompressed))
641 return None, None
642
643 strides = [0] * len(augmented_shape)
644 stride = self.element_size() * self.storage_compression_scale
645
646 if self.format != TensorFormat.NHCWB16:
647 for i in stride_order:
648 strides[i] = stride
649 stride *= augmented_shape[i]
650 else:
651 assert len(strides) == 5
Tim Hall79d07d22020-04-27 18:20:16 +0100652 strides[4] = stride
Patrik Gustavsson2213e902020-05-05 17:49:35 +0200653 strides[3] = 16 * stride # STRIDE_X
Tim Hall79d07d22020-04-27 18:20:16 +0100654 strides[1] = strides[3] * augmented_shape[2] # STRIDE_C
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200655 strides[2] = augmented_shape[2] * augmented_shape[3] * stride # STRIDE_Y
Tim Hall79d07d22020-04-27 18:20:16 +0100656 strides[0] = strides[2] * augmented_shape[1] # STRIDE_N
657
658 return strides, augmented_coord
659
660 def get_strides(self):
661 strides, _ = self.get_strides_and_coord()
662
663 return strides
664
Louis Verhaard3c07c972020-05-07 08:12:58 +0200665 def needs_dma(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200666 return len(self.ops) == 1 and self.ops[0].type == Op.DMA
Louis Verhaard3c07c972020-05-07 08:12:58 +0200667
668 def get_dma_src_tensor(self):
669 # For weight tensors that need DMA: returns the source tensor in Flash, else None
670 # Note: for DMA ops, Pass.weight_tensor is referring to the SRAM weight tensor
671 return self.ops[0].inputs[0] if self.needs_dma() else None
672
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200673 def find_npu_op(self):
674 # Returns the NPU operator that uses this tensor, excluding DMA operators.
675 for op in self.consumers():
Louis Verhaardaee5d752020-09-30 09:01:52 +0200676 if op.type == Op.DMA:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200677 return op.outputs[0].find_npu_op()
Dwight Lidman940fdee2020-08-13 13:11:48 +0200678 if op.run_on_npu:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200679 return op
680 return None
681
Tim Hall79d07d22020-04-27 18:20:16 +0100682 def compressed_stream_index_from_coord(self, coord):
683 assert self.format == TensorFormat.WeightsCompressed
684 assert len(self.compressed_values) > 0
685 assert len(self.compressed_values) + 1 == len(self.weight_compressed_offsets)
686
687 depth = coord[-1]
688 brick_depth = self.brick_size[-1]
689 # Clamp position at final element index
690 if depth > self.shape[-1]:
691 depth = self.shape[-1]
692
693 # Always round up to next boundary
Michael McGeagh8d3216f2020-08-10 11:35:57 +0100694 index = numeric_util.round_up_divide(depth, brick_depth)
Tim Hall79d07d22020-04-27 18:20:16 +0100695
696 # Check boundaries on all but last weight set (which may be shorter
697 # than the brick we divided it up into)
698 if index < len(self.weight_compressed_offsets) - 1:
699 # There are no half-way points in the weights
700 if (depth % brick_depth) != 0:
701 raise Exception("Offset into weights must be aligned to a brick")
702
703 return index
704
705 def size_of_compressed_stream(self, index):
706 assert 0 <= index < len(self.compressed_values)
707 return len(self.compressed_values[index])
708
709 def is_last_index_in_compressed_stream(self, index):
710 assert 0 <= index < len(self.compressed_values)
711 return index == len(self.compressed_values) - 1
712
713 def address_offset_for_coordinate(self, orig_coord, is_top_box=False):
714 address_offset = 0
715 coord = orig_coord
716
717 coord = coord[-len(self.storage_shape) :]
718
719 if self.sub_purpose == TensorSubPurpose.Standard:
720 for idx, c in enumerate(coord):
721 if is_top_box:
722 assert c > 0 and c <= self.shape[idx]
723 else:
724 assert c >= 0 and c < self.shape[idx]
725
726 if self.format == TensorFormat.WeightsCompressed:
727 if len(self.weight_compressed_offsets) == 0:
728 return 0
729
Louis Verhaard3c07c972020-05-07 08:12:58 +0200730 if self.needs_dma() and self.sub_purpose == TensorSubPurpose.DoubleBuffer:
Tim Hall79d07d22020-04-27 18:20:16 +0100731 depth = orig_coord[-1]
732 brick_depth = self.brick_size[-1]
733 # Clamp position at final element index
734 if depth > self.shape[-1]:
735 depth = self.shape[-1]
736
737 # Always round up to next boundary
Michael McGeagh8d3216f2020-08-10 11:35:57 +0100738 index = numeric_util.round_up_divide(depth, brick_depth)
Tim Hall79d07d22020-04-27 18:20:16 +0100739 index = index % 2
740
741 if len(self.compressed_values) <= 2:
742 if is_top_box and index == 0:
743 for cv in self.compressed_values:
744 address_offset += len(cv)
745 else:
746 address_offset = index * len(self.compressed_values[0])
747 else:
748 if is_top_box and index == 0:
749 address_offset = self.storage_shape[-1]
750 else:
751 address_offset = index * (self.storage_shape[-1] // 2)
752 else:
753 index = self.compressed_stream_index_from_coord(orig_coord)
754 assert index < len(self.weight_compressed_offsets)
755 address_offset = self.weight_compressed_offsets[index]
756 else:
757 if is_top_box:
758 coord = [c - 1 for c in coord]
759
760 # handle wraparound for partial buffers. make sure to do this after subtracting top box:
761 coord = [c % self.storage_shape[idx] for idx, c in enumerate(coord)]
762
763 strides, augmented_coord = self.get_strides_and_coord(coord)
764 if strides is None:
765 return None
766
767 if is_top_box:
768 address_offset += 1 * strides[-1] # one element
769
770 address_offset += np.dot(augmented_coord, strides)
771
772 assert address_offset >= 0
773 assert address_offset <= self.storage_size()
774 return address_offset
775
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200776 def is_allocated_in_tensor_arena(self, scratch_tensor_mem_area):
777 if self.mem_area == scratch_tensor_mem_area and (self.mem_type in set((MemType.Scratch, MemType.Scratch_fast))):
778 return True
779 return False
780
Louis Verhaard0b8268a2020-08-05 16:11:29 +0200781 def equivalent(self, tens):
782 return self.equivalence_id == tens.equivalence_id
783
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100784 def set_all_shapes(self, shape):
785 self.shape = shape
786 self.storage_shape = shape
787 self.bandwidth_shape = shape
788
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100789 def get_full_shape(self):
790 d = len(self.shape)
791 if d in (1, 3):
Michael McGeagh8d3216f2020-08-10 11:35:57 +0100792 return numeric_util.full_shape(4, self.shape, 1)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100793 elif d == 2:
794 return [self.shape[0], 1, 1, self.shape[1]]
795 else:
Fredrik Svedberg835d8e12020-09-04 09:46:17 +0200796 return self.shape.copy()
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100797
Tim Hall93582962020-09-09 21:58:15 +0100798 def is_quantized(self):
799 # a tensor is quantized if it has an integral type and it contains valid quantization params
800
Tim Hall89567612020-10-27 11:57:57 +0000801 if not isinstance(self.quantization, QuantizationParameters):
Tim Hall93582962020-09-09 21:58:15 +0100802 return False
803
Tim Hall89567612020-10-27 11:57:57 +0000804 return (self.dtype.type & BaseType.Int) != 0 and self.quantization.is_valid()
Tim Hall93582962020-09-09 21:58:15 +0100805
Tim Hall79d07d22020-04-27 18:20:16 +0100806 def __str__(self):
807 return "<nng.Tensor '%s' shape=%s dtype=%s>" % (self.name, self.shape, self.dtype)
808
809 __repr__ = __str__
Tim Hall93582962020-09-09 21:58:15 +0100810
811
Tim Hall93582962020-09-09 21:58:15 +0100812def check_quantized_tens_scaling_equal(tens_a, tens_b):
813 # checks that the scaling of two quantized tensors are equal
814
Tim Hall89567612020-10-27 11:57:57 +0000815 return tens_a.is_quantized() and tens_b.is_quantized() and tens_a.quantization.is_scaling_equal(tens_b.quantization)