blob: 97885d0e88bf31655eeaaf07bd9c420d878c0dc8 [file] [log] [blame]
erik.andersson@arm.com42b94ed2021-02-11 14:02:08 +01001# Copyright (C) 2020-2021 Arm Limited or its affiliates. All rights reserved.
Tim Hall79d07d22020-04-27 18:20:16 +01002#
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
Diqing Zhongf842b692020-12-11 13:07:37 +010022from enum import auto
Louis Verhaard9db529a2020-09-23 10:27:11 +020023from functools import lru_cache
Louis Verhaard6c74c3b2020-12-17 13:54:09 +010024from functools import total_ordering
Louis Verhaard93719a92020-12-08 10:02:31 +010025from typing import Dict
26from typing import List
27from typing import Optional
28from typing import Tuple
29from typing import Union
30from uuid import UUID
Diego Russoea6111a2020-04-14 18:41:58 +010031
32import numpy as np
33
34from . import numeric_util
Tim Hall93582962020-09-09 21:58:15 +010035from .data_type import BaseType
Michael McGeagh5778ffd2020-08-06 17:31:02 +010036from .data_type import DataType
Michael McGeagh528a56d2020-12-16 11:33:21 +000037from .errors import UnsupportedFeatureError
38from .errors import VelaError
Dwight Lidmana9390f72020-05-13 12:00:08 +020039from .ethos_u55_regs.ethos_u55_regs import resampling_mode
Patrik Gustavsson2349d422020-12-01 16:02:29 +010040from .numeric_util import full_shape
Louis Verhaardaee5d752020-09-30 09:01:52 +020041from .operation import Op
Michael McGeagh5778ffd2020-08-06 17:31:02 +010042from .operation import Operation
patrik.gustavssoneeb85152020-12-21 17:10:40 +000043from .shape4d import Shape4D
Louis Verhaard93719a92020-12-08 10:02:31 +010044
45Shape = List
Tim Hall79d07d22020-04-27 18:20:16 +010046
47
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020048class MemType(enum.IntFlag):
49 Unknown = 0
50 Permanent_NPU = 1
51 Permanent_CPU = 2
52 Scratch = 3
53 Scratch_fast = 4
54 Size = Scratch_fast + 1
55
Louis Verhaard93719a92020-12-08 10:02:31 +010056 def display_name(self) -> str:
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020057 return ("Unknown", "Permanent_NPU", "Permanent_CPU", "Scratch", "Scratch_fast", "Size")[self.value]
58
Louis Verhaard93719a92020-12-08 10:02:31 +010059 def identifier_name(self) -> str:
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020060 return ("unknown", "permanent_npu", "permanent_cpu", "scratch", "scratch_fast", "size")[self.value]
61
Louis Verhaard93719a92020-12-08 10:02:31 +010062 @staticmethod
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020063 def all():
64 return (MemType.Permanent_NPU, MemType.Permanent_CPU, MemType.Scratch, MemType.Scratch_fast)
65
66 def __str__(self):
67 return self.name
68
69
Diqing Zhongf842b692020-12-11 13:07:37 +010070class BandwidthDirection(enum.IntEnum):
71 Read = 0
72 Write = auto()
73 Size = auto()
74
75 def display_name(self):
76 return self.name
77
78 def identifier_name(self):
79 return self.name.lower()
80
81 @staticmethod
82 def all():
83 return (BandwidthDirection.Read, BandwidthDirection.Write)
84
85
Tim Hall79d07d22020-04-27 18:20:16 +010086class MemArea(enum.IntFlag):
87 Unknown = 0
88 Sram = 1
89 Dram = 2
90 OnChipFlash = 3
91 OffChipFlash = 4
Louis Verhaard0b8268a2020-08-05 16:11:29 +020092 Shram = 5 # for LUT
93 Size = Shram + 1
Tim Hall79d07d22020-04-27 18:20:16 +010094
Louis Verhaard93719a92020-12-08 10:02:31 +010095 def display_name(self) -> str:
Louis Verhaard0b8268a2020-08-05 16:11:29 +020096 return ("Unknown", "SRAM", "DRAM", "On-chip Flash", "Off-chip Flash", "SHRAM", "Size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +010097
Louis Verhaard93719a92020-12-08 10:02:31 +010098 def identifier_name(self) -> str:
Louis Verhaard0b8268a2020-08-05 16:11:29 +020099 return ("unknown", "sram", "dram", "on_chip_flash", "off_chip_flash", "shram", "size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +0100100
Louis Verhaard93719a92020-12-08 10:02:31 +0100101 @staticmethod
Tim Hall79d07d22020-04-27 18:20:16 +0100102 def all():
Louis Verhaard0b8268a2020-08-05 16:11:29 +0200103 return (MemArea.Sram, MemArea.Dram, MemArea.OnChipFlash, MemArea.OffChipFlash, MemArea.Shram)
Tim Hall79d07d22020-04-27 18:20:16 +0100104
105 def __str__(self):
106 return self.name
107
108
109class TensorPurpose(enum.IntFlag):
110 Unknown = 0
111 Weights = 1
112 FeatureMap = 2
113 Scratch = 3
Fredrik Svedberge22ba8c2021-01-27 16:53:41 +0100114 ScratchFast = 4
115 LUT = 5
116 FSBias = 6
117 Size = 7
Tim Hall79d07d22020-04-27 18:20:16 +0100118
Louis Verhaard93719a92020-12-08 10:02:31 +0100119 def display_name(self) -> str:
Fredrik Svedberge22ba8c2021-01-27 16:53:41 +0100120 return ("Unknown", "Weights", "FeatureMap", "Scratch", "ScratchFast", "LUT", "FastStorageBias", "Size")[
121 self.value
122 ]
Tim Hall79d07d22020-04-27 18:20:16 +0100123
Louis Verhaard93719a92020-12-08 10:02:31 +0100124 def identifier_name(self) -> str:
Fredrik Svedberge22ba8c2021-01-27 16:53:41 +0100125 return ("unknown", "weights", "feature_map", "scratch", "scratch_fast", "lut", "fast_storage_bias", "size")[
126 self.value
127 ]
Tim Hall79d07d22020-04-27 18:20:16 +0100128
Louis Verhaard93719a92020-12-08 10:02:31 +0100129 @staticmethod
Tim Hall79d07d22020-04-27 18:20:16 +0100130 def all():
Andreas Nevalainen897cc142020-10-28 15:42:08 +0100131 return (TensorPurpose.Weights, TensorPurpose.FeatureMap, TensorPurpose.FSBias)
Tim Hall79d07d22020-04-27 18:20:16 +0100132
133
134class TensorSubPurpose(enum.Enum):
135 Standard = 0
136 DoubleBuffer = 1
137 RollingBufferX = 2
138 RollingBufferY = 3
139 RollingBufferXY = 4
140
Louis Verhaard93719a92020-12-08 10:02:31 +0100141 def display_name(self) -> str:
Tim Hall79d07d22020-04-27 18:20:16 +0100142 return ("Standard", "Double Buffer", "Rolling Buffer X", "Rolling Buffer Y", "Rolling Buffer XY")[self.value]
143
Louis Verhaard93719a92020-12-08 10:02:31 +0100144 def identifier_name(self) -> str:
Tim Hall79d07d22020-04-27 18:20:16 +0100145 return ("standard", "double_buffer", "rolling_buffer_x", "rolling_buffer_y", "rolling_buffer_xy")[self.value]
146
Louis Verhaard93719a92020-12-08 10:02:31 +0100147 @staticmethod
Tim Hall79d07d22020-04-27 18:20:16 +0100148 def all():
149 return (
150 TensorSubPurpose.Standard,
151 TensorSubPurpose.DoubleBuffer,
152 TensorSubPurpose.RollingBufferX,
153 TensorSubPurpose.RollingBufferY,
154 TensorSubPurpose.RollingBufferXY,
155 )
156
157
158class TensorFormat(enum.Flag):
159 Unknown = 0
160 WeightsCompressed = 1
161 NHWC = 2
162 NHCWB16 = 3
163
164 def __str__(self):
165 return self.name
166
167
168class TensorBlockTraversal(enum.Enum):
169 Default = 0
170 DepthWise = 1
171 DepthFirst = 2
172 PartKernelFirst = 3
173
174
Louis Verhaard93719a92020-12-08 10:02:31 +0100175def shape_num_elements(shp: Shape) -> Optional[int]:
Tim Hall79d07d22020-04-27 18:20:16 +0100176 elems = 1
177 if shp is None:
178 return None
179 for d in shp:
180 if d is None:
181 return None
182 elems *= d
183 return elems
184
185
Louis Verhaard93719a92020-12-08 10:02:31 +0100186def shape_fully_defined(shp: Shape) -> bool:
Tim Hall79d07d22020-04-27 18:20:16 +0100187 if shp is None:
188 return False
189 for d in shp:
190 if d is None:
191 return False
192 return True
193
194
Louis Verhaard93719a92020-12-08 10:02:31 +0100195def shape_round_to_quantum(shp: Shape, quantum: Tuple) -> Shape:
Tim Hall79d07d22020-04-27 18:20:16 +0100196 new_shp = list(shp)
197
198 # Traverse backwards using length of shape since there may be more rounding quantums than shape elements
199 for i in range(-1, -len(shp) - 1, -1):
200 if new_shp[i] is not None:
201 new_shp[i] = numeric_util.round_up(new_shp[i], quantum[i])
202 return new_shp
203
204
Louis Verhaard9db529a2020-09-23 10:27:11 +0200205@lru_cache(maxsize=None)
Louis Verhaard93719a92020-12-08 10:02:31 +0100206def create_equivalence_id(key) -> UUID:
Louis Verhaard9db529a2020-09-23 10:27:11 +0200207 # Generates equivalence_id based on the given key.
208 return uuid.uuid4()
209
210
Tim Hall79d07d22020-04-27 18:20:16 +0100211class QuantizationParameters:
212 __slots__ = "min", "max", "num_bits", "narrow_range", "scale_f32", "zero_point", "quant_min", "quant_max"
213
Louis Verhaard93719a92020-12-08 10:02:31 +0100214 def __init__(
215 self,
216 min: Union[float, np.ndarray, None] = None,
217 max: Union[float, np.ndarray, None] = None,
218 num_bits=None,
219 narrow_range=None,
220 ):
Tim Hall79d07d22020-04-27 18:20:16 +0100221 self.min = min
222 self.max = max
223
224 self.num_bits = num_bits
225 self.narrow_range = narrow_range
226
Louis Verhaard93719a92020-12-08 10:02:31 +0100227 self.scale_f32: Union[float, np.ndarray, None] = None
228 self.zero_point: Union[int, np.ndarray, None] = None
229 self.quant_min: Optional[float] = None
230 self.quant_max: Optional[float] = None
Tim Hall79d07d22020-04-27 18:20:16 +0100231
232 def __str__(self):
233 return "<nng.QuantizationParameters min=%s max=%s, num_bits=%s, scale=%s, zero_point=%s>" % (
234 self.min,
235 self.max,
236 self.num_bits,
237 self.scale_f32,
238 self.zero_point,
239 )
240
241 __repr__ = __str__
242
Louis Verhaard93719a92020-12-08 10:02:31 +0100243 def clone(self) -> "QuantizationParameters":
Tim Hall79d07d22020-04-27 18:20:16 +0100244 res = QuantizationParameters()
245 res.min = self.min
246 res.max = self.max
247
248 res.num_bits = self.num_bits
249 res.narrow_range = self.narrow_range
250
251 res.scale_f32 = self.scale_f32
252 res.zero_point = self.zero_point
253 res.quant_min = self.quant_min
254 res.quant_max = self.quant_max
255 return res
256
257 def dequantize(self, values):
258 if self.zero_point.size == 1 and self.scale_f32.size == 1:
259 # same scale is used for all values
260 res = (values.astype(np.float64) - self.zero_point) * self.scale_f32
261 else:
262 # a different scale is used for different sets of values
263 values_as_float = values.astype(np.float64)
264
265 # this is not compatible with the format of depthwise weights,
266 # where input is at index 3 (Output, Kh, Kw, Input)
267 # return the quantized values
268 return np.ndarray((values_as_float.shape))
269
Tim Hall79d07d22020-04-27 18:20:16 +0100270 return res
271
Louis Verhaard93719a92020-12-08 10:02:31 +0100272 def is_scaling_equal(self, other: Optional["QuantizationParameters"]) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100273 # quantisation parameter scaling is not equal if 'other' is None because
274 # it implies that the tensor it belongs to is not quantised. otherwise,
275 # it depends upon whether the scale and zero point are equal
276
Tim Hall89567612020-10-27 11:57:57 +0000277 if not isinstance(other, QuantizationParameters):
Tim Halle3786ac2020-07-28 17:40:50 +0100278 return False
279
280 return self.scale_f32 == other.scale_f32 and self.zero_point == other.zero_point
281
Louis Verhaard93719a92020-12-08 10:02:31 +0100282 def is_valid(self) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100283 # quantisation parameters are consider valid if they have a scale and zero point
284
285 return None not in (self.scale_f32, self.zero_point)
286
Louis Verhaard93719a92020-12-08 10:02:31 +0100287 def is_per_axis(self) -> bool:
Dwight Lidmanc7187432020-11-16 17:40:46 +0100288 """Returns True if either the scale, zero point, minimum or maximum values are arrays"""
289 for attr in ("scale_f32", "zero_point", "min", "max"):
290 if isinstance(getattr(self, attr), np.ndarray):
291 return True
292 return False
293
Tim Hall79d07d22020-04-27 18:20:16 +0100294
Louis Verhaard93719a92020-12-08 10:02:31 +0100295def create_const_tensor(
296 name: str,
297 shape: Shape,
298 dtype: DataType,
299 values: np.ndarray,
300 value_dtype: np.dtype = None,
301 purpose: TensorPurpose = TensorPurpose.Unknown,
302 quantization: QuantizationParameters = None,
303):
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100304 # Tensor
305 const_tensor = Tensor(shape, dtype, name + "_0")
306 const_tensor.purpose = purpose
307 const_tensor.quantization = quantization
308 const_tensor.values = np.array(values, dtype=value_dtype)
Jacob Bohlina41cd4d2020-08-26 18:21:28 +0200309 const_tensor.quant_values = np.frombuffer(const_tensor.values.tobytes(), dtype=np.uint8)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100310 # Operator
Louis Verhaardaee5d752020-09-30 09:01:52 +0200311 const_op = Operation(Op.Const, name)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100312 const_op.set_output_tensor(const_tensor)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000313 const_op.set_ifm_ofm_shapes()
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100314 return const_tensor
315
316
Jacob Bohlin1a666972020-09-11 10:04:15 +0200317# class that keeps track of all tensor addresses in the different memory types
318class TensorAddressMap:
Louis Verhaard93719a92020-12-08 10:02:31 +0100319 address_map: Dict = defaultdict(dict) # dict (tens.equivalence_id -> dict (mem_type -> address))
Jacob Bohlin1a666972020-09-11 10:04:15 +0200320
321 @classmethod
Louis Verhaard93719a92020-12-08 10:02:31 +0100322 def get_address_for_tens(cls, tens_id: UUID, mem_type: MemType) -> int:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200323 return cls.address_map[tens_id].get(mem_type)
324
325 @classmethod
Louis Verhaard93719a92020-12-08 10:02:31 +0100326 def set_address_for_tens(cls, tens_id: UUID, mem_type: MemType, address: int):
Jacob Bohlin1a666972020-09-11 10:04:15 +0200327 # Check previous address if there is one
328 previous_address = cls.address_map[tens_id].get(mem_type)
Louis Verhaard0b9c9a32020-09-15 14:05:38 +0200329 if address is not None and previous_address is not None:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200330 assert previous_address == address, "Two different addresses cannot be assigned to the same tensor."
331
332 # Set tensor's address for memory type
333 cls.address_map[tens_id][mem_type] = address
334
335
Louis Verhaard6c74c3b2020-12-17 13:54:09 +0100336@total_ordering
Tim Hall79d07d22020-04-27 18:20:16 +0100337class Tensor:
338 __slots__ = (
339 "shape",
340 "storage_shape",
341 "bandwidth_shape",
342 "dtype",
343 "name",
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100344 "is_variable",
Tim Hall79d07d22020-04-27 18:20:16 +0100345 "ops",
346 "consumer_list",
347 "values",
348 "quant_values",
349 "compressed_values",
Tim Hallf7e810a2020-06-25 15:04:31 +0100350 "compressed_values_substream_offsets",
Tim Hall79d07d22020-04-27 18:20:16 +0100351 "mem_area",
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200352 "mem_type",
Tim Hall79d07d22020-04-27 18:20:16 +0100353 "format",
354 "purpose",
355 "sub_purpose",
356 "alignment",
357 "weight_transpose_depthwise",
358 "storage_compression_scale",
359 "bandwidth_compression_scale",
360 "compression_scale_for_worst_weight_stream",
361 "weight_compression_scales",
362 "weight_compression_config",
Louis Verhaard9db529a2020-09-23 10:27:11 +0200363 "value_id",
Tim Hall79d07d22020-04-27 18:20:16 +0100364 "storage_rounding_quantum",
365 "brick_size",
Tim Hall79d07d22020-04-27 18:20:16 +0100366 "quantization",
367 "weight_compressed_offsets",
368 "element_size_bytes",
Tim Hall79d07d22020-04-27 18:20:16 +0100369 "block_traversal",
Tim Hall79d07d22020-04-27 18:20:16 +0100370 "equivalence_id",
Dwight Lidmana9390f72020-05-13 12:00:08 +0200371 "resampling_mode",
Patrik Gustavsson458a2082020-08-13 13:41:05 +0200372 "avoid_NHCWB16",
Tim Hall79d07d22020-04-27 18:20:16 +0100373 )
374 AllocationQuantum = 16
375
Louis Verhaard93719a92020-12-08 10:02:31 +0100376 def __init__(self, shape: Shape, dtype: DataType, name: str):
Tim Hall79d07d22020-04-27 18:20:16 +0100377 self.shape = shape
378 self.storage_shape = shape
379 self.bandwidth_shape = shape
380 self.dtype = dtype
381 self.name = name
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100382 self.is_variable = False
Louis Verhaard93719a92020-12-08 10:02:31 +0100383 self.equivalence_id: UUID = uuid.uuid4()
Tim Hall79d07d22020-04-27 18:20:16 +0100384
Louis Verhaard93719a92020-12-08 10:02:31 +0100385 self.ops: List[Operation] = []
386 self.consumer_list: List[Operation] = []
Tim Hall79d07d22020-04-27 18:20:16 +0100387
Louis Verhaard93719a92020-12-08 10:02:31 +0100388 self.values: Optional[np.ndarray] = None
389 self.quant_values: Optional[np.ndarray] = None
390 self.compressed_values: Optional[np.ndarray] = None
391 self.compressed_values_substream_offsets: Optional[List] = None
392 self.mem_area: MemArea = MemArea.Unknown
393 self.mem_type: MemType = MemType.Unknown
394 self.format: TensorFormat = TensorFormat.Unknown
395 self.purpose: TensorPurpose = TensorPurpose.Unknown
396 self.sub_purpose: TensorSubPurpose = TensorSubPurpose.Standard
397 self.alignment: int = Tensor.AllocationQuantum
398 self.weight_transpose_depthwise: bool = False
Tim Hall79d07d22020-04-27 18:20:16 +0100399
Louis Verhaard93719a92020-12-08 10:02:31 +0100400 self.storage_compression_scale: float = 1.0
401 self.bandwidth_compression_scale: float = 1.0
402 self.compression_scale_for_worst_weight_stream: float = 1.0
403 self.weight_compression_scales: Optional[np.ndarray] = None
Louis Verhaard9db529a2020-09-23 10:27:11 +0200404 # if two tensors have the same weight_compression_config, then they have the same compressed values
Tim Hall79d07d22020-04-27 18:20:16 +0100405 self.weight_compression_config = None
Louis Verhaard9db529a2020-09-23 10:27:11 +0200406 # if two tensors have the same value_id, then they have the same values
Louis Verhaard93719a92020-12-08 10:02:31 +0100407 self.value_id: UUID = uuid.uuid4()
408 self.weight_compressed_offsets: List = []
409 self.storage_rounding_quantum: Tuple = (1, 1, 1, 1)
410 self.brick_size: Tuple = (1, 1, 1, 1)
411 self.element_size_bytes: int = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100412
413 # quantization parameters
Louis Verhaard93719a92020-12-08 10:02:31 +0100414 self.quantization: Optional[QuantizationParameters] = None
415 self.block_traversal: TensorBlockTraversal = TensorBlockTraversal.Default
416 self.resampling_mode: resampling_mode = resampling_mode.NONE
Tim Hall79d07d22020-04-27 18:20:16 +0100417
Louis Verhaard93719a92020-12-08 10:02:31 +0100418 self.avoid_NHCWB16: bool = False
Patrik Gustavsson458a2082020-08-13 13:41:05 +0200419
Jacob Bohlin1a666972020-09-11 10:04:15 +0200420 @property
Louis Verhaard93719a92020-12-08 10:02:31 +0100421 def address(self) -> int:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200422 return TensorAddressMap.get_address_for_tens(self.equivalence_id, self.mem_type)
423
424 @address.setter
Louis Verhaard93719a92020-12-08 10:02:31 +0100425 def address(self, address: int):
Jacob Bohlin1a666972020-09-11 10:04:15 +0200426 TensorAddressMap.set_address_for_tens(self.equivalence_id, self.mem_type, address)
427
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100428 @property
429 def is_standard_fm(self) -> bool:
430 return self.sub_purpose == TensorSubPurpose.Standard and self.purpose == TensorPurpose.FeatureMap
431
Louis Verhaard93719a92020-12-08 10:02:31 +0100432 def element_size(self) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100433 if self.element_size_bytes == 0:
434 return self.dtype.size_in_bits() / 8
435 return self.element_size_bytes
436
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100437 # Returns a copy, renamed to self.name + suffix
438 # The references to Operators will be empty when returned
439 # Depending on set_unique, the copy is shallow, or deep
440 # For set_unique==True, a new equivalence_id will be set
Louis Verhaard93719a92020-12-08 10:02:31 +0100441 def clone(self, suffix="_clone", set_unique: bool = False) -> "Tensor":
erik.andersson@arm.com42b94ed2021-02-11 14:02:08 +0100442 res = copy.copy(self)
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100443 if set_unique:
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100444 res.equivalence_id = uuid.uuid4()
erik.andersson@arm.com42b94ed2021-02-11 14:02:08 +0100445 res.storage_shape = list(self.storage_shape)
446 res.bandwidth_shape = list(self.bandwidth_shape)
447 if self.quantization is not None:
448 res.quantization = self.quantization.clone()
Tim Hall79d07d22020-04-27 18:20:16 +0100449
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100450 res.name = res.name + suffix
Tim Hall79d07d22020-04-27 18:20:16 +0100451 res.ops = []
452 res.consumer_list = []
Tim Hall79d07d22020-04-27 18:20:16 +0100453
Tim Hall79d07d22020-04-27 18:20:16 +0100454 return res
455
Louis Verhaard93719a92020-12-08 10:02:31 +0100456 def clone_into_fast_storage(self, arch) -> "Tensor":
Tim Hall79d07d22020-04-27 18:20:16 +0100457 res = self.clone(suffix="_fast_storage")
458 res.mem_area = arch.fast_storage_mem_area
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200459 res.mem_type = MemType.Scratch_fast
Tim Hall79d07d22020-04-27 18:20:16 +0100460 return res
461
Louis Verhaard93719a92020-12-08 10:02:31 +0100462 def copy_compressed_weight_info(self, src_tens: "Tensor"):
Louis Verhaard3c07c972020-05-07 08:12:58 +0200463 # Copies compressed values + all related weight compression info from the given tensor
Louis Verhaard9db529a2020-09-23 10:27:11 +0200464 self.equivalence_id = src_tens.equivalence_id
Louis Verhaard3c07c972020-05-07 08:12:58 +0200465 self.compressed_values = src_tens.compressed_values
Tim Hallf7e810a2020-06-25 15:04:31 +0100466 self.compressed_values_substream_offsets = src_tens.compressed_values_substream_offsets
Louis Verhaard3c07c972020-05-07 08:12:58 +0200467 self.storage_shape = src_tens.storage_shape
468 self.brick_size = src_tens.brick_size
469 self.weight_compression_scales = src_tens.weight_compression_scales
470 self.weight_compressed_offsets = src_tens.weight_compressed_offsets
471 self.weight_transpose_depthwise = src_tens.weight_transpose_depthwise
472 self.compression_scale_for_worst_weight_stream = src_tens.compression_scale_for_worst_weight_stream
473 self.storage_compression_scale = src_tens.storage_compression_scale
Diqing Zhong7e1d1d12020-10-30 15:10:46 +0100474 self.bandwidth_compression_scale = src_tens.bandwidth_compression_scale
Louis Verhaard3c07c972020-05-07 08:12:58 +0200475 self.block_traversal = src_tens.block_traversal
476 self.weight_compression_config = src_tens.weight_compression_config
Louis Verhaard9db529a2020-09-23 10:27:11 +0200477 self.value_id = src_tens.value_id
Louis Verhaard3c07c972020-05-07 08:12:58 +0200478
Louis Verhaard93719a92020-12-08 10:02:31 +0100479 def set_format(self, fmt: TensorFormat, arch):
Tim Hall79d07d22020-04-27 18:20:16 +0100480 self.format = fmt
481 shape_len = 0
482 try:
483 shape_len = len(self.shape)
484 except TypeError:
485 pass
486
Louis Verhaard0411edb2020-11-16 16:37:11 +0100487 if shape_len > 4:
488 return
Tim Hall79d07d22020-04-27 18:20:16 +0100489 self.storage_rounding_quantum = arch.storage_rounding_quantums[self.format]
Louis Verhaard93719a92020-12-08 10:02:31 +0100490 self.storage_rounding_quantum = tuple(self.storage_rounding_quantum[-shape_len:])
Tim Hall79d07d22020-04-27 18:20:16 +0100491 self.brick_size = arch.brick_sizes[self.format]
Louis Verhaard93719a92020-12-08 10:02:31 +0100492 self.brick_size = tuple(self.brick_size[-shape_len:])
Tim Hall79d07d22020-04-27 18:20:16 +0100493 if self.shape is None:
494 return
495
496 self.bandwidth_shape = shape_round_to_quantum(self.shape, self.brick_size)
497 self.storage_shape = shape_round_to_quantum(self.shape, self.storage_rounding_quantum)
498
499 if fmt == TensorFormat.WeightsCompressed:
500 compression_ratio = 5 / 8
501 self.storage_compression_scale = compression_ratio
502 self.bandwidth_compression_scale = compression_ratio
503 self.compression_scale_for_worst_weight_stream = compression_ratio
504
Louis Verhaard93719a92020-12-08 10:02:31 +0100505 def storage_elements(self) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100506 elems = shape_num_elements(self.storage_shape)
507 if elems is None:
508 return 0
509 return elems
510
Louis Verhaard93719a92020-12-08 10:02:31 +0100511 def elements(self) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100512 elems = shape_num_elements(self.shape)
513 if elems is None:
514 return 0
515 return elems
516
Louis Verhaard93719a92020-12-08 10:02:31 +0100517 def has_fully_defined_shape(self) -> bool:
Tim Hall79d07d22020-04-27 18:20:16 +0100518 return shape_fully_defined(self.shape)
519
Louis Verhaard93719a92020-12-08 10:02:31 +0100520 def storage_size(self, scale: float = 1.0) -> int:
Patrik Gustavsson90831bc2020-08-24 16:26:11 +0200521 raw_size = self.storage_elements() * self.element_size() * scale
Tim Hall79d07d22020-04-27 18:20:16 +0100522 if raw_size == 0:
523 raw_size = 1 # force it to take up space
524 rounded_size = numeric_util.round_up(numeric_util.round_up_to_int(raw_size), self.alignment)
525 return rounded_size
526
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100527 def storage_size_for_shape(self, op_storage_shape: Shape) -> int:
528 elems = shape_num_elements(op_storage_shape)
529 elems = elems if elems else 0
530 raw_size = elems * self.element_size()
531 if raw_size == 0:
532 raw_size = 1 # force it to take up space
533 rounded_size = numeric_util.round_up(numeric_util.round_up_to_int(raw_size), self.alignment)
534 return rounded_size
535
Louis Verhaard93719a92020-12-08 10:02:31 +0100536 def storage_size_for_sub_purpose(
537 self, arch, sub_purpose: TensorSubPurpose, param_a: Optional[int] = None, param_b: Optional[int] = None
538 ) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100539 alt_shape = self.storage_shape_for_sub_purpose(sub_purpose, param_a, param_b)
540 elems = shape_num_elements(alt_shape)
541 if elems is None:
542 return 0
543 if sub_purpose == TensorSubPurpose.DoubleBuffer:
Patrik Gustavsson90831bc2020-08-24 16:26:11 +0200544 raw_size = (
545 elems
546 * self.element_size()
547 * self.compression_scale_for_worst_weight_stream
548 * arch.weight_estimation_scaling
549 )
Tim Hall79d07d22020-04-27 18:20:16 +0100550 else:
Patrik Gustavsson9baa4c32020-08-20 13:59:01 +0200551 # Rolling buffers are used for intermediate data in ifm streaming
552 # These will all use the NHCWB16 format, and need to be aligned to 16 in the C-dimension
553 if alt_shape[-1] % 16 != 0:
554 nhcwb16_shape = alt_shape[0:-1] + [numeric_util.round_up(alt_shape[-1], 16)]
555 elems = shape_num_elements(nhcwb16_shape)
556
Tim Hall79d07d22020-04-27 18:20:16 +0100557 raw_size = elems * self.element_size() * self.storage_compression_scale
558 rounded_size = numeric_util.round_up(numeric_util.round_up_to_int(raw_size), self.alignment)
559 return rounded_size
560
Louis Verhaard93719a92020-12-08 10:02:31 +0100561 def storage_shape_for_sub_purpose(
562 self, sub_purpose: TensorSubPurpose, param_a: Optional[int], param_b: Optional[int]
563 ) -> Shape:
Tim Hall79d07d22020-04-27 18:20:16 +0100564 if sub_purpose == TensorSubPurpose.DoubleBuffer:
Jacob Bohline843d332020-06-23 12:12:56 +0200565 shp = list(self.shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100566 assert len(shp) >= 2
Louis Verhaard93719a92020-12-08 10:02:31 +0100567 assert param_a is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100568 shp[-1] = min(shp[-1], param_a * 2)
Tim Hall79d07d22020-04-27 18:20:16 +0100569 else:
Jacob Bohline843d332020-06-23 12:12:56 +0200570 shp = list(self.storage_shape)
571 if sub_purpose == TensorSubPurpose.RollingBufferX:
572 assert len(shp) == 4
Louis Verhaard93719a92020-12-08 10:02:31 +0100573 assert param_a is not None
Jacob Bohline843d332020-06-23 12:12:56 +0200574 shp[0] = 1
575 shp[2] = min(shp[2], param_a)
576 elif sub_purpose == TensorSubPurpose.RollingBufferY:
577 assert len(shp) == 4
Louis Verhaard93719a92020-12-08 10:02:31 +0100578 assert param_a is not None
Jacob Bohline843d332020-06-23 12:12:56 +0200579 shp[0] = 1
580 shp[1] = min(shp[1], param_a)
581 elif sub_purpose == TensorSubPurpose.RollingBufferXY:
582 assert len(shp) == 4
Louis Verhaard93719a92020-12-08 10:02:31 +0100583 assert param_a is not None
584 assert param_b is not None
Jacob Bohline843d332020-06-23 12:12:56 +0200585 shp[0] = 1
586 shp[2] = min(shp[2], param_a)
587 shp[1] = min(shp[1], param_b)
588 elif sub_purpose == TensorSubPurpose.Standard:
589 pass
590 else:
591 assert 0, "did not expect new sub purpose %s" % (sub_purpose,)
592
Tim Hall79d07d22020-04-27 18:20:16 +0100593 return shp
594
Louis Verhaard93719a92020-12-08 10:02:31 +0100595 def set_new_sub_purpose(self, sub_purpose: TensorSubPurpose, param_a=None, param_b=None):
Tim Hall79d07d22020-04-27 18:20:16 +0100596 self.storage_shape = self.storage_shape_for_sub_purpose(sub_purpose, param_a, param_b)
597 self.sub_purpose = sub_purpose
598 if sub_purpose == TensorSubPurpose.DoubleBuffer:
599 self.storage_compression_scale = self.compression_scale_for_worst_weight_stream
600
Louis Verhaard93719a92020-12-08 10:02:31 +0100601 def bandwidth(self) -> float:
Tim Hall79d07d22020-04-27 18:20:16 +0100602 elems = shape_num_elements(self.bandwidth_shape)
603 if elems is None:
604 return 0
605 return elems * self.element_size() * self.bandwidth_compression_scale
606
Louis Verhaard93719a92020-12-08 10:02:31 +0100607 def consumers(self) -> List[Operation]:
Tim Hall79d07d22020-04-27 18:20:16 +0100608 return self.consumer_list
609
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100610 def get_4D_storage_shape_for_shape(self, op_shape4D: Shape4D) -> Shape4D:
611 rounding_quantum = full_shape(4, list(self.storage_rounding_quantum), 1)
612 return Shape4D(shape_round_to_quantum(op_shape4D.as_list(), rounding_quantum))
613
614 def addresses_for_rolling_buffer(self, start_coord: Shape, end_coord: Shape, op_shape4D: Shape4D) -> Tuple:
Tim Hall79d07d22020-04-27 18:20:16 +0100615 # returns ( box_height0, box_height1, box_width, [address_tl, address_tr, address_bl, address_br] )
616
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100617 if self.storage_shape == []:
618 return (
619 1,
620 1,
621 1,
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100622 [self.address_for_coordinate(start_coord, op_shape4D=op_shape4D), None, None, None],
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100623 )
Tim Hall79d07d22020-04-27 18:20:16 +0100624
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100625 if self.is_standard_fm:
626 storage_shape_4D = self.get_4D_storage_shape_for_shape(op_shape4D)
627 else:
628 storage_shape_4D = Shape4D(self.storage_shape)
629
630 crossing_y = numeric_util.round_up(start_coord[1] + 1, storage_shape_4D.height)
631 crossing_x = numeric_util.round_up(start_coord[2] + 1, storage_shape_4D.width)
Tim Hall79d07d22020-04-27 18:20:16 +0100632
633 crossing_y = min(crossing_y, end_coord[1])
634 crossing_x = min(crossing_x, end_coord[2])
635
636 box_height0 = crossing_y - start_coord[1]
637 box_width = crossing_x - start_coord[2]
638
Louis Verhaard93719a92020-12-08 10:02:31 +0100639 addresses: List = [None] * 4
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100640 addresses[0] = self.address_for_coordinate(start_coord, op_shape4D=op_shape4D)
Tim Hall79d07d22020-04-27 18:20:16 +0100641
642 if end_coord[2] > crossing_x:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100643 addresses[1] = self.address_for_coordinate(
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100644 [start_coord[0], start_coord[1], crossing_x, start_coord[3]], op_shape4D=op_shape4D
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100645 )
Michael McGeagh528a56d2020-12-16 11:33:21 +0000646 raise UnsupportedFeatureError("Striping in vertical direction is not supported")
Tim Hall79d07d22020-04-27 18:20:16 +0100647 if end_coord[1] > crossing_y:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100648 addresses[2] = self.address_for_coordinate(
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100649 [start_coord[0], crossing_y, start_coord[2], start_coord[3]], op_shape4D=op_shape4D
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100650 )
Tim Hall79d07d22020-04-27 18:20:16 +0100651 if end_coord[1] > crossing_y and end_coord[2] > crossing_x:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100652 addresses[3] = self.address_for_coordinate(
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100653 [start_coord[0], crossing_y, crossing_x, start_coord[3]], op_shape4D=op_shape4D
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100654 )
Tim Hall79d07d22020-04-27 18:20:16 +0100655
656 return box_height0, box_height0, box_width, addresses
657
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100658 def address_for_coordinate(self, coord: Shape, is_top_box: bool = False, op_shape4D: Shape4D = None) -> int:
659 offset = self.address_offset_for_coordinate(coord, op_shape4D=op_shape4D, is_top_box=is_top_box)
Louis Verhaard93719a92020-12-08 10:02:31 +0100660 assert offset is not None
661 return self.address + offset
Tim Hall79d07d22020-04-27 18:20:16 +0100662
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100663 def get_strides_and_coord(
664 self, coord: Optional[Shape] = None, shape4D: Optional[Shape4D] = None
665 ) -> Tuple[Optional[Shape], Optional[Shape]]:
Tim Hall79d07d22020-04-27 18:20:16 +0100666 if coord is None:
667 coord = [0] * len(self.storage_shape)
668
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100669 if shape4D and self.is_standard_fm:
670 augmented_shape = self.get_4D_storage_shape_for_shape(shape4D).as_list()
671 else:
672 augmented_shape = full_shape(4, self.storage_shape, 1)
673
Tim Hall79d07d22020-04-27 18:20:16 +0100674 augmented_coord = coord
Tim Hall79d07d22020-04-27 18:20:16 +0100675
676 while len(augmented_coord) < 4:
677 augmented_coord = [0] + augmented_coord
678
679 assert len(augmented_coord) == len(augmented_shape)
680
681 if self.format == TensorFormat.NHWC:
682 augmented_shape = [augmented_shape[0], augmented_shape[3]] + augmented_shape[1:3] + [1]
683 augmented_coord = [augmented_coord[0], augmented_coord[3]] + augmented_coord[1:3] + [0]
Tim Hall79d07d22020-04-27 18:20:16 +0100684
685 elif self.format == TensorFormat.NHCWB16:
Patrik Gustavsson2213e902020-05-05 17:49:35 +0200686 channel_divisor = 16
Tim Hall79d07d22020-04-27 18:20:16 +0100687 augmented_shape = augmented_shape[0:4] + [1]
688 augmented_coord = (
689 [augmented_coord[0], augmented_coord[3] // channel_divisor]
690 + augmented_coord[1:3]
691 + [augmented_coord[3] % channel_divisor]
692 )
693
694 if augmented_shape[1] == 0:
695 augmented_shape[1] = 1
696
697 else:
Michael McGeaghf3e3ad72020-12-02 12:39:03 +0000698 assert self.format in (TensorFormat.Unknown, TensorFormat.WeightsCompressed)
Tim Hall79d07d22020-04-27 18:20:16 +0100699 return None, None
700
Louis Verhaard93719a92020-12-08 10:02:31 +0100701 strides: List = [0] * len(augmented_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100702 stride = self.element_size() * self.storage_compression_scale
703
704 if self.format != TensorFormat.NHCWB16:
Louis Verhaard93719a92020-12-08 10:02:31 +0100705 stride_order = [4, 1, 3, 2, 0]
Tim Hall79d07d22020-04-27 18:20:16 +0100706 for i in stride_order:
707 strides[i] = stride
708 stride *= augmented_shape[i]
709 else:
710 assert len(strides) == 5
Tim Hall79d07d22020-04-27 18:20:16 +0100711 strides[4] = stride
Patrik Gustavsson2213e902020-05-05 17:49:35 +0200712 strides[3] = 16 * stride # STRIDE_X
Tim Hall79d07d22020-04-27 18:20:16 +0100713 strides[1] = strides[3] * augmented_shape[2] # STRIDE_C
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200714 strides[2] = augmented_shape[2] * augmented_shape[3] * stride # STRIDE_Y
Tim Hall79d07d22020-04-27 18:20:16 +0100715 strides[0] = strides[2] * augmented_shape[1] # STRIDE_N
716
717 return strides, augmented_coord
718
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100719 def get_strides(self, shape4D: Optional[Shape4D] = None) -> Shape:
720 strides, _ = self.get_strides_and_coord(shape4D=shape4D)
Louis Verhaard93719a92020-12-08 10:02:31 +0100721 assert strides is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100722 return strides
723
Louis Verhaard93719a92020-12-08 10:02:31 +0100724 def needs_dma(self) -> bool:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200725 return len(self.ops) == 1 and self.ops[0].type == Op.DMA
Louis Verhaard3c07c972020-05-07 08:12:58 +0200726
Louis Verhaard93719a92020-12-08 10:02:31 +0100727 def get_dma_src_tensor(self) -> "Optional[Tensor]":
Louis Verhaard3c07c972020-05-07 08:12:58 +0200728 # For weight tensors that need DMA: returns the source tensor in Flash, else None
729 # Note: for DMA ops, Pass.weight_tensor is referring to the SRAM weight tensor
730 return self.ops[0].inputs[0] if self.needs_dma() else None
731
Louis Verhaard93719a92020-12-08 10:02:31 +0100732 def find_npu_op(self) -> Optional[Operation]:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200733 # Returns the NPU operator that uses this tensor, excluding DMA operators.
734 for op in self.consumers():
Louis Verhaardaee5d752020-09-30 09:01:52 +0200735 if op.type == Op.DMA:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200736 return op.outputs[0].find_npu_op()
Dwight Lidman940fdee2020-08-13 13:11:48 +0200737 if op.run_on_npu:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200738 return op
Louis Verhaard93719a92020-12-08 10:02:31 +0100739 return None
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200740
Louis Verhaard93719a92020-12-08 10:02:31 +0100741 def compressed_stream_index_from_coord(self, coord: Shape) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100742 assert self.format == TensorFormat.WeightsCompressed
Louis Verhaard93719a92020-12-08 10:02:31 +0100743 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100744 assert len(self.compressed_values) > 0
745 assert len(self.compressed_values) + 1 == len(self.weight_compressed_offsets)
746
747 depth = coord[-1]
748 brick_depth = self.brick_size[-1]
749 # Clamp position at final element index
750 if depth > self.shape[-1]:
751 depth = self.shape[-1]
752
753 # Always round up to next boundary
Michael McGeagh8d3216f2020-08-10 11:35:57 +0100754 index = numeric_util.round_up_divide(depth, brick_depth)
Tim Hall79d07d22020-04-27 18:20:16 +0100755
756 # Check boundaries on all but last weight set (which may be shorter
757 # than the brick we divided it up into)
758 if index < len(self.weight_compressed_offsets) - 1:
759 # There are no half-way points in the weights
760 if (depth % brick_depth) != 0:
Michael McGeagh528a56d2020-12-16 11:33:21 +0000761 raise UnsupportedFeatureError("Offset into weights must be aligned to a brick")
Tim Hall79d07d22020-04-27 18:20:16 +0100762
763 return index
764
Louis Verhaard93719a92020-12-08 10:02:31 +0100765 def size_of_compressed_stream(self, index: int) -> int:
766 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100767 assert 0 <= index < len(self.compressed_values)
768 return len(self.compressed_values[index])
769
Louis Verhaard93719a92020-12-08 10:02:31 +0100770 def is_last_index_in_compressed_stream(self, index: int) -> bool:
771 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100772 assert 0 <= index < len(self.compressed_values)
773 return index == len(self.compressed_values) - 1
774
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100775 def address_offset_for_coordinate(
776 self, orig_coord: Shape, op_shape4D: Optional[Shape4D] = None, is_top_box: bool = False
777 ) -> Optional[int]:
Tim Hall79d07d22020-04-27 18:20:16 +0100778 address_offset = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100779
780 if self.sub_purpose == TensorSubPurpose.Standard:
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100781 shape = op_shape4D.as_list() if op_shape4D else self.shape
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100782 for idx, c in enumerate(orig_coord):
Tim Hall79d07d22020-04-27 18:20:16 +0100783 if is_top_box:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100784 assert c > 0 and c <= shape[idx]
Tim Hall79d07d22020-04-27 18:20:16 +0100785 else:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100786 assert c >= 0 and c < shape[idx]
Tim Hall79d07d22020-04-27 18:20:16 +0100787
788 if self.format == TensorFormat.WeightsCompressed:
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100789 storage_size = self.storage_size()
Tim Hall79d07d22020-04-27 18:20:16 +0100790 if len(self.weight_compressed_offsets) == 0:
791 return 0
792
Louis Verhaard3c07c972020-05-07 08:12:58 +0200793 if self.needs_dma() and self.sub_purpose == TensorSubPurpose.DoubleBuffer:
Tim Hall79d07d22020-04-27 18:20:16 +0100794 depth = orig_coord[-1]
795 brick_depth = self.brick_size[-1]
796 # Clamp position at final element index
797 if depth > self.shape[-1]:
798 depth = self.shape[-1]
799
800 # Always round up to next boundary
Michael McGeagh8d3216f2020-08-10 11:35:57 +0100801 index = numeric_util.round_up_divide(depth, brick_depth)
Tim Hall79d07d22020-04-27 18:20:16 +0100802 index = index % 2
Louis Verhaard93719a92020-12-08 10:02:31 +0100803 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100804
805 if len(self.compressed_values) <= 2:
806 if is_top_box and index == 0:
807 for cv in self.compressed_values:
808 address_offset += len(cv)
809 else:
810 address_offset = index * len(self.compressed_values[0])
811 else:
812 if is_top_box and index == 0:
813 address_offset = self.storage_shape[-1]
814 else:
815 address_offset = index * (self.storage_shape[-1] // 2)
816 else:
817 index = self.compressed_stream_index_from_coord(orig_coord)
818 assert index < len(self.weight_compressed_offsets)
819 address_offset = self.weight_compressed_offsets[index]
820 else:
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100821 coord = orig_coord
822 if op_shape4D and self.is_standard_fm:
823 storage_shape = self.get_4D_storage_shape_for_shape(op_shape4D).as_list()
824 storage_size = self.storage_size_for_shape(storage_shape)
825 else:
826 storage_shape = self.storage_shape
827 coord = coord[-len(storage_shape) :]
828 storage_size = self.storage_size()
829
Tim Hall79d07d22020-04-27 18:20:16 +0100830 if is_top_box:
831 coord = [c - 1 for c in coord]
832
833 # handle wraparound for partial buffers. make sure to do this after subtracting top box:
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100834 coord = [c % storage_shape[idx] for idx, c in enumerate(coord)]
Tim Hall79d07d22020-04-27 18:20:16 +0100835
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100836 strides, augmented_coord = self.get_strides_and_coord(coord, op_shape4D)
Tim Hall79d07d22020-04-27 18:20:16 +0100837 if strides is None:
838 return None
839
840 if is_top_box:
841 address_offset += 1 * strides[-1] # one element
842
843 address_offset += np.dot(augmented_coord, strides)
844
845 assert address_offset >= 0
Patrik Gustavsson3a269202021-01-21 08:28:55 +0100846 assert address_offset <= storage_size
Tim Hall79d07d22020-04-27 18:20:16 +0100847 return address_offset
848
Louis Verhaard93719a92020-12-08 10:02:31 +0100849 def is_allocated_in_tensor_arena(self, scratch_tensor_mem_area: MemArea) -> bool:
Michael McGeaghf3e3ad72020-12-02 12:39:03 +0000850 return (self.mem_area == scratch_tensor_mem_area) and (self.mem_type in (MemType.Scratch, MemType.Scratch_fast))
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200851
Louis Verhaard93719a92020-12-08 10:02:31 +0100852 def equivalent(self, tens: "Tensor") -> bool:
Louis Verhaard0b8268a2020-08-05 16:11:29 +0200853 return self.equivalence_id == tens.equivalence_id
854
Louis Verhaard93719a92020-12-08 10:02:31 +0100855 def set_all_shapes(self, shape: Shape):
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100856 self.shape = shape
857 self.storage_shape = shape
858 self.bandwidth_shape = shape
859
Louis Verhaard93719a92020-12-08 10:02:31 +0100860 def get_full_shape(self) -> Shape:
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100861 d = len(self.shape)
862 if d in (1, 3):
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100863 return full_shape(4, self.shape, 1)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100864 elif d == 2:
865 return [self.shape[0], 1, 1, self.shape[1]]
866 else:
Fredrik Svedberg835d8e12020-09-04 09:46:17 +0200867 return self.shape.copy()
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100868
Louis Verhaard93719a92020-12-08 10:02:31 +0100869 def is_quantized(self) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100870 # a tensor is quantized if it has an integral type and it contains valid quantization params
871
Tim Hall89567612020-10-27 11:57:57 +0000872 if not isinstance(self.quantization, QuantizationParameters):
Tim Hall93582962020-09-09 21:58:15 +0100873 return False
874
Tim Hall89567612020-10-27 11:57:57 +0000875 return (self.dtype.type & BaseType.Int) != 0 and self.quantization.is_valid()
Tim Hall93582962020-09-09 21:58:15 +0100876
Louis Verhaard6c74c3b2020-12-17 13:54:09 +0100877 def __lt__(self, other: "Tensor") -> bool:
878 return self.equivalence_id < other.equivalence_id
879
Tim Hall79d07d22020-04-27 18:20:16 +0100880 def __str__(self):
881 return "<nng.Tensor '%s' shape=%s dtype=%s>" % (self.name, self.shape, self.dtype)
882
883 __repr__ = __str__
Tim Hall93582962020-09-09 21:58:15 +0100884
Michael McGeagh528a56d2020-12-16 11:33:21 +0000885 def error(self, msg):
886 """
887 Raises a VelaError exception for errors encountered when parsing a Tensor
888
889 :param self: Tensor object that resulted in the error
890 :param msg: str object that contains a description of the specific error encountered
891 """
892
893 def _print_operators(ops):
894 lines = []
895 for idx, op in enumerate(ops):
896 op_type = getattr(op, "type", "Not an Operation")
897 op_id = getattr(op, "op_index", "-")
898 lines.append(f" {idx} = {op_type} ({op_id})")
899 return lines
900
901 lines = [f"Invalid {self.name} tensor. {msg}"]
902
903 lines += [" Driving operators:"]
904 lines += _print_operators(self.ops)
905
906 lines += [" Consuming operators:"]
907 lines += _print_operators(self.consumer_list)
908
909 raise VelaError("\n".join(lines))
910
Tim Hall93582962020-09-09 21:58:15 +0100911
Louis Verhaard93719a92020-12-08 10:02:31 +0100912def check_quantized_tens_scaling_equal(tens_a: Tensor, tens_b: Tensor) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100913 # checks that the scaling of two quantized tensors are equal
914
Tim Hall89567612020-10-27 11:57:57 +0000915 return tens_a.is_quantized() and tens_b.is_quantized() and tens_a.quantization.is_scaling_equal(tens_b.quantization)