blob: 093e8771bbeabc3c9407aa292b4d6681bd9a386a [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
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 Svedberga0c36242020-06-03 15:43:31 +0200114 LUT = 4
Andreas Nevalainen897cc142020-10-28 15:42:08 +0100115 FSBias = 5
116 Size = 6
Tim Hall79d07d22020-04-27 18:20:16 +0100117
Louis Verhaard93719a92020-12-08 10:02:31 +0100118 def display_name(self) -> str:
Andreas Nevalainen897cc142020-10-28 15:42:08 +0100119 return ("Unknown", "Weights", "FeatureMap", "Scratch", "LUT", "FastStorageBias", "Size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +0100120
Louis Verhaard93719a92020-12-08 10:02:31 +0100121 def identifier_name(self) -> str:
Andreas Nevalainen897cc142020-10-28 15:42:08 +0100122 return ("unknown", "weights", "feature_map", "scratch", "lut", "fast_storage_bias", "size")[self.value]
Tim Hall79d07d22020-04-27 18:20:16 +0100123
Louis Verhaard93719a92020-12-08 10:02:31 +0100124 @staticmethod
Tim Hall79d07d22020-04-27 18:20:16 +0100125 def all():
Andreas Nevalainen897cc142020-10-28 15:42:08 +0100126 return (TensorPurpose.Weights, TensorPurpose.FeatureMap, TensorPurpose.FSBias)
Tim Hall79d07d22020-04-27 18:20:16 +0100127
128
129class TensorSubPurpose(enum.Enum):
130 Standard = 0
131 DoubleBuffer = 1
132 RollingBufferX = 2
133 RollingBufferY = 3
134 RollingBufferXY = 4
135
Louis Verhaard93719a92020-12-08 10:02:31 +0100136 def display_name(self) -> str:
Tim Hall79d07d22020-04-27 18:20:16 +0100137 return ("Standard", "Double Buffer", "Rolling Buffer X", "Rolling Buffer Y", "Rolling Buffer XY")[self.value]
138
Louis Verhaard93719a92020-12-08 10:02:31 +0100139 def identifier_name(self) -> str:
Tim Hall79d07d22020-04-27 18:20:16 +0100140 return ("standard", "double_buffer", "rolling_buffer_x", "rolling_buffer_y", "rolling_buffer_xy")[self.value]
141
Louis Verhaard93719a92020-12-08 10:02:31 +0100142 @staticmethod
Tim Hall79d07d22020-04-27 18:20:16 +0100143 def all():
144 return (
145 TensorSubPurpose.Standard,
146 TensorSubPurpose.DoubleBuffer,
147 TensorSubPurpose.RollingBufferX,
148 TensorSubPurpose.RollingBufferY,
149 TensorSubPurpose.RollingBufferXY,
150 )
151
152
153class TensorFormat(enum.Flag):
154 Unknown = 0
155 WeightsCompressed = 1
156 NHWC = 2
157 NHCWB16 = 3
158
159 def __str__(self):
160 return self.name
161
162
163class TensorBlockTraversal(enum.Enum):
164 Default = 0
165 DepthWise = 1
166 DepthFirst = 2
167 PartKernelFirst = 3
168
169
Louis Verhaard93719a92020-12-08 10:02:31 +0100170def shape_num_elements(shp: Shape) -> Optional[int]:
Tim Hall79d07d22020-04-27 18:20:16 +0100171 elems = 1
172 if shp is None:
173 return None
174 for d in shp:
175 if d is None:
176 return None
177 elems *= d
178 return elems
179
180
Louis Verhaard93719a92020-12-08 10:02:31 +0100181def shape_fully_defined(shp: Shape) -> bool:
Tim Hall79d07d22020-04-27 18:20:16 +0100182 if shp is None:
183 return False
184 for d in shp:
185 if d is None:
186 return False
187 return True
188
189
Louis Verhaard93719a92020-12-08 10:02:31 +0100190def shape_round_to_quantum(shp: Shape, quantum: Tuple) -> Shape:
Tim Hall79d07d22020-04-27 18:20:16 +0100191 new_shp = list(shp)
192
193 # Traverse backwards using length of shape since there may be more rounding quantums than shape elements
194 for i in range(-1, -len(shp) - 1, -1):
195 if new_shp[i] is not None:
196 new_shp[i] = numeric_util.round_up(new_shp[i], quantum[i])
197 return new_shp
198
199
Louis Verhaard9db529a2020-09-23 10:27:11 +0200200@lru_cache(maxsize=None)
Louis Verhaard93719a92020-12-08 10:02:31 +0100201def create_equivalence_id(key) -> UUID:
Louis Verhaard9db529a2020-09-23 10:27:11 +0200202 # Generates equivalence_id based on the given key.
203 return uuid.uuid4()
204
205
Tim Hall79d07d22020-04-27 18:20:16 +0100206class QuantizationParameters:
207 __slots__ = "min", "max", "num_bits", "narrow_range", "scale_f32", "zero_point", "quant_min", "quant_max"
208
Louis Verhaard93719a92020-12-08 10:02:31 +0100209 def __init__(
210 self,
211 min: Union[float, np.ndarray, None] = None,
212 max: Union[float, np.ndarray, None] = None,
213 num_bits=None,
214 narrow_range=None,
215 ):
Tim Hall79d07d22020-04-27 18:20:16 +0100216 self.min = min
217 self.max = max
218
219 self.num_bits = num_bits
220 self.narrow_range = narrow_range
221
Louis Verhaard93719a92020-12-08 10:02:31 +0100222 self.scale_f32: Union[float, np.ndarray, None] = None
223 self.zero_point: Union[int, np.ndarray, None] = None
224 self.quant_min: Optional[float] = None
225 self.quant_max: Optional[float] = None
Tim Hall79d07d22020-04-27 18:20:16 +0100226
227 def __str__(self):
228 return "<nng.QuantizationParameters min=%s max=%s, num_bits=%s, scale=%s, zero_point=%s>" % (
229 self.min,
230 self.max,
231 self.num_bits,
232 self.scale_f32,
233 self.zero_point,
234 )
235
236 __repr__ = __str__
237
Louis Verhaard93719a92020-12-08 10:02:31 +0100238 def clone(self) -> "QuantizationParameters":
Tim Hall79d07d22020-04-27 18:20:16 +0100239 res = QuantizationParameters()
240 res.min = self.min
241 res.max = self.max
242
243 res.num_bits = self.num_bits
244 res.narrow_range = self.narrow_range
245
246 res.scale_f32 = self.scale_f32
247 res.zero_point = self.zero_point
248 res.quant_min = self.quant_min
249 res.quant_max = self.quant_max
250 return res
251
252 def dequantize(self, values):
253 if self.zero_point.size == 1 and self.scale_f32.size == 1:
254 # same scale is used for all values
255 res = (values.astype(np.float64) - self.zero_point) * self.scale_f32
256 else:
257 # a different scale is used for different sets of values
258 values_as_float = values.astype(np.float64)
259
260 # this is not compatible with the format of depthwise weights,
261 # where input is at index 3 (Output, Kh, Kw, Input)
262 # return the quantized values
263 return np.ndarray((values_as_float.shape))
264
Tim Hall79d07d22020-04-27 18:20:16 +0100265 return res
266
Louis Verhaard93719a92020-12-08 10:02:31 +0100267 def is_scaling_equal(self, other: Optional["QuantizationParameters"]) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100268 # quantisation parameter scaling is not equal if 'other' is None because
269 # it implies that the tensor it belongs to is not quantised. otherwise,
270 # it depends upon whether the scale and zero point are equal
271
Tim Hall89567612020-10-27 11:57:57 +0000272 if not isinstance(other, QuantizationParameters):
Tim Halle3786ac2020-07-28 17:40:50 +0100273 return False
274
275 return self.scale_f32 == other.scale_f32 and self.zero_point == other.zero_point
276
Louis Verhaard93719a92020-12-08 10:02:31 +0100277 def is_valid(self) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100278 # quantisation parameters are consider valid if they have a scale and zero point
279
280 return None not in (self.scale_f32, self.zero_point)
281
Louis Verhaard93719a92020-12-08 10:02:31 +0100282 def is_per_axis(self) -> bool:
Dwight Lidmanc7187432020-11-16 17:40:46 +0100283 """Returns True if either the scale, zero point, minimum or maximum values are arrays"""
284 for attr in ("scale_f32", "zero_point", "min", "max"):
285 if isinstance(getattr(self, attr), np.ndarray):
286 return True
287 return False
288
Tim Hall79d07d22020-04-27 18:20:16 +0100289
Louis Verhaard93719a92020-12-08 10:02:31 +0100290def create_const_tensor(
291 name: str,
292 shape: Shape,
293 dtype: DataType,
294 values: np.ndarray,
295 value_dtype: np.dtype = None,
296 purpose: TensorPurpose = TensorPurpose.Unknown,
297 quantization: QuantizationParameters = None,
298):
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100299 # Tensor
300 const_tensor = Tensor(shape, dtype, name + "_0")
301 const_tensor.purpose = purpose
302 const_tensor.quantization = quantization
303 const_tensor.values = np.array(values, dtype=value_dtype)
Jacob Bohlina41cd4d2020-08-26 18:21:28 +0200304 const_tensor.quant_values = np.frombuffer(const_tensor.values.tobytes(), dtype=np.uint8)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100305 # Operator
Louis Verhaardaee5d752020-09-30 09:01:52 +0200306 const_op = Operation(Op.Const, name)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100307 const_op.set_output_tensor(const_tensor)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000308 const_op.set_ifm_ofm_shapes()
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100309 return const_tensor
310
311
312def create_reshape_tensor(tens, shape, ifm_reshape=True):
313 if shape == tens.shape:
314 return tens
315 # Tensors
316 name = tens.name + "_reshape"
317 reshape_ifm = tens
318 reshape_ofm = tens.clone("_reshaped")
319 reshape_ofm.set_all_shapes(shape)
320 if not ifm_reshape:
321 reshape_ifm, reshape_ofm = reshape_ofm, reshape_ifm
322 # Operator
Louis Verhaardaee5d752020-09-30 09:01:52 +0200323 reshape_op = Operation(Op.Reshape, name)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100324 reshape_op.attrs["new_shape"] = shape
325 reshape_op.add_input_tensor(reshape_ifm)
326 reshape_op.add_input_tensor(create_const_tensor(name + "_shape", [1], DataType.int32, shape))
327 reshape_op.set_output_tensor(reshape_ofm)
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000328 reshape_op.set_ifm_ofm_shapes()
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100329 return reshape_ofm if ifm_reshape else reshape_ifm
330
331
Jacob Bohlin1a666972020-09-11 10:04:15 +0200332# class that keeps track of all tensor addresses in the different memory types
333class TensorAddressMap:
Louis Verhaard93719a92020-12-08 10:02:31 +0100334 address_map: Dict = defaultdict(dict) # dict (tens.equivalence_id -> dict (mem_type -> address))
Jacob Bohlin1a666972020-09-11 10:04:15 +0200335
336 @classmethod
Louis Verhaard93719a92020-12-08 10:02:31 +0100337 def get_address_for_tens(cls, tens_id: UUID, mem_type: MemType) -> int:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200338 return cls.address_map[tens_id].get(mem_type)
339
340 @classmethod
Louis Verhaard93719a92020-12-08 10:02:31 +0100341 def set_address_for_tens(cls, tens_id: UUID, mem_type: MemType, address: int):
Jacob Bohlin1a666972020-09-11 10:04:15 +0200342 # Check previous address if there is one
343 previous_address = cls.address_map[tens_id].get(mem_type)
Louis Verhaard0b9c9a32020-09-15 14:05:38 +0200344 if address is not None and previous_address is not None:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200345 assert previous_address == address, "Two different addresses cannot be assigned to the same tensor."
346
347 # Set tensor's address for memory type
348 cls.address_map[tens_id][mem_type] = address
349
350
Louis Verhaard6c74c3b2020-12-17 13:54:09 +0100351@total_ordering
Tim Hall79d07d22020-04-27 18:20:16 +0100352class Tensor:
353 __slots__ = (
354 "shape",
355 "storage_shape",
356 "bandwidth_shape",
357 "dtype",
358 "name",
359 "ops",
360 "consumer_list",
361 "values",
362 "quant_values",
363 "compressed_values",
Tim Hallf7e810a2020-06-25 15:04:31 +0100364 "compressed_values_substream_offsets",
Tim Hall79d07d22020-04-27 18:20:16 +0100365 "mem_area",
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200366 "mem_type",
Tim Hall79d07d22020-04-27 18:20:16 +0100367 "format",
368 "purpose",
369 "sub_purpose",
370 "alignment",
371 "weight_transpose_depthwise",
372 "storage_compression_scale",
373 "bandwidth_compression_scale",
374 "compression_scale_for_worst_weight_stream",
375 "weight_compression_scales",
376 "weight_compression_config",
Louis Verhaard9db529a2020-09-23 10:27:11 +0200377 "value_id",
Tim Hall79d07d22020-04-27 18:20:16 +0100378 "storage_rounding_quantum",
379 "brick_size",
Tim Hall79d07d22020-04-27 18:20:16 +0100380 "quantization",
381 "weight_compressed_offsets",
382 "element_size_bytes",
Tim Hall79d07d22020-04-27 18:20:16 +0100383 "block_traversal",
Tim Hall79d07d22020-04-27 18:20:16 +0100384 "equivalence_id",
Dwight Lidmana9390f72020-05-13 12:00:08 +0200385 "resampling_mode",
Patrik Gustavsson458a2082020-08-13 13:41:05 +0200386 "avoid_NHCWB16",
Tim Hall79d07d22020-04-27 18:20:16 +0100387 )
388 AllocationQuantum = 16
389
Louis Verhaard93719a92020-12-08 10:02:31 +0100390 def __init__(self, shape: Shape, dtype: DataType, name: str):
Tim Hall79d07d22020-04-27 18:20:16 +0100391 self.shape = shape
392 self.storage_shape = shape
393 self.bandwidth_shape = shape
394 self.dtype = dtype
395 self.name = name
Louis Verhaard93719a92020-12-08 10:02:31 +0100396 self.equivalence_id: UUID = uuid.uuid4()
Tim Hall79d07d22020-04-27 18:20:16 +0100397
Louis Verhaard93719a92020-12-08 10:02:31 +0100398 self.ops: List[Operation] = []
399 self.consumer_list: List[Operation] = []
Tim Hall79d07d22020-04-27 18:20:16 +0100400
Louis Verhaard93719a92020-12-08 10:02:31 +0100401 self.values: Optional[np.ndarray] = None
402 self.quant_values: Optional[np.ndarray] = None
403 self.compressed_values: Optional[np.ndarray] = None
404 self.compressed_values_substream_offsets: Optional[List] = None
405 self.mem_area: MemArea = MemArea.Unknown
406 self.mem_type: MemType = MemType.Unknown
407 self.format: TensorFormat = TensorFormat.Unknown
408 self.purpose: TensorPurpose = TensorPurpose.Unknown
409 self.sub_purpose: TensorSubPurpose = TensorSubPurpose.Standard
410 self.alignment: int = Tensor.AllocationQuantum
411 self.weight_transpose_depthwise: bool = False
Tim Hall79d07d22020-04-27 18:20:16 +0100412
Louis Verhaard93719a92020-12-08 10:02:31 +0100413 self.storage_compression_scale: float = 1.0
414 self.bandwidth_compression_scale: float = 1.0
415 self.compression_scale_for_worst_weight_stream: float = 1.0
416 self.weight_compression_scales: Optional[np.ndarray] = None
Louis Verhaard9db529a2020-09-23 10:27:11 +0200417 # if two tensors have the same weight_compression_config, then they have the same compressed values
Tim Hall79d07d22020-04-27 18:20:16 +0100418 self.weight_compression_config = None
Louis Verhaard9db529a2020-09-23 10:27:11 +0200419 # if two tensors have the same value_id, then they have the same values
Louis Verhaard93719a92020-12-08 10:02:31 +0100420 self.value_id: UUID = uuid.uuid4()
421 self.weight_compressed_offsets: List = []
422 self.storage_rounding_quantum: Tuple = (1, 1, 1, 1)
423 self.brick_size: Tuple = (1, 1, 1, 1)
424 self.element_size_bytes: int = 0
Tim Hall79d07d22020-04-27 18:20:16 +0100425
426 # quantization parameters
Louis Verhaard93719a92020-12-08 10:02:31 +0100427 self.quantization: Optional[QuantizationParameters] = None
428 self.block_traversal: TensorBlockTraversal = TensorBlockTraversal.Default
429 self.resampling_mode: resampling_mode = resampling_mode.NONE
Tim Hall79d07d22020-04-27 18:20:16 +0100430
Louis Verhaard93719a92020-12-08 10:02:31 +0100431 self.avoid_NHCWB16: bool = False
Patrik Gustavsson458a2082020-08-13 13:41:05 +0200432
Jacob Bohlin1a666972020-09-11 10:04:15 +0200433 @property
Louis Verhaard93719a92020-12-08 10:02:31 +0100434 def address(self) -> int:
Jacob Bohlin1a666972020-09-11 10:04:15 +0200435 return TensorAddressMap.get_address_for_tens(self.equivalence_id, self.mem_type)
436
437 @address.setter
Louis Verhaard93719a92020-12-08 10:02:31 +0100438 def address(self, address: int):
Jacob Bohlin1a666972020-09-11 10:04:15 +0200439 TensorAddressMap.set_address_for_tens(self.equivalence_id, self.mem_type, address)
440
Louis Verhaard93719a92020-12-08 10:02:31 +0100441 def element_size(self) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100442 if self.element_size_bytes == 0:
443 return self.dtype.size_in_bits() / 8
444 return self.element_size_bytes
445
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100446 # Returns a copy, renamed to self.name + suffix
447 # The references to Operators will be empty when returned
448 # Depending on set_unique, the copy is shallow, or deep
449 # For set_unique==True, a new equivalence_id will be set
Louis Verhaard93719a92020-12-08 10:02:31 +0100450 def clone(self, suffix="_clone", set_unique: bool = False) -> "Tensor":
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100451 if set_unique:
452 res = copy.deepcopy(self)
453 res.equivalence_id = uuid.uuid4()
454 else:
455 res = copy.copy(self)
456 res.storage_shape = list(self.storage_shape)
457 res.bandwidth_shape = list(self.bandwidth_shape)
458 if self.quantization is not None:
459 res.quantization = self.quantization.clone()
Tim Hall79d07d22020-04-27 18:20:16 +0100460
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100461 res.name = res.name + suffix
Tim Hall79d07d22020-04-27 18:20:16 +0100462 res.ops = []
463 res.consumer_list = []
Tim Hall79d07d22020-04-27 18:20:16 +0100464
Tim Hall79d07d22020-04-27 18:20:16 +0100465 return res
466
Louis Verhaard93719a92020-12-08 10:02:31 +0100467 def clone_into_fast_storage(self, arch) -> "Tensor":
Tim Hall79d07d22020-04-27 18:20:16 +0100468 res = self.clone(suffix="_fast_storage")
469 res.mem_area = arch.fast_storage_mem_area
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200470 res.mem_type = MemType.Scratch_fast
Tim Hall79d07d22020-04-27 18:20:16 +0100471 return res
472
Louis Verhaard93719a92020-12-08 10:02:31 +0100473 def copy_compressed_weight_info(self, src_tens: "Tensor"):
Louis Verhaard3c07c972020-05-07 08:12:58 +0200474 # Copies compressed values + all related weight compression info from the given tensor
Louis Verhaard9db529a2020-09-23 10:27:11 +0200475 self.equivalence_id = src_tens.equivalence_id
Louis Verhaard3c07c972020-05-07 08:12:58 +0200476 self.compressed_values = src_tens.compressed_values
Tim Hallf7e810a2020-06-25 15:04:31 +0100477 self.compressed_values_substream_offsets = src_tens.compressed_values_substream_offsets
Louis Verhaard3c07c972020-05-07 08:12:58 +0200478 self.storage_shape = src_tens.storage_shape
479 self.brick_size = src_tens.brick_size
480 self.weight_compression_scales = src_tens.weight_compression_scales
481 self.weight_compressed_offsets = src_tens.weight_compressed_offsets
482 self.weight_transpose_depthwise = src_tens.weight_transpose_depthwise
483 self.compression_scale_for_worst_weight_stream = src_tens.compression_scale_for_worst_weight_stream
484 self.storage_compression_scale = src_tens.storage_compression_scale
Diqing Zhong7e1d1d12020-10-30 15:10:46 +0100485 self.bandwidth_compression_scale = src_tens.bandwidth_compression_scale
Louis Verhaard3c07c972020-05-07 08:12:58 +0200486 self.block_traversal = src_tens.block_traversal
487 self.weight_compression_config = src_tens.weight_compression_config
Louis Verhaard9db529a2020-09-23 10:27:11 +0200488 self.value_id = src_tens.value_id
Louis Verhaard3c07c972020-05-07 08:12:58 +0200489
Louis Verhaard93719a92020-12-08 10:02:31 +0100490 def set_format(self, fmt: TensorFormat, arch):
Tim Hall79d07d22020-04-27 18:20:16 +0100491 self.format = fmt
492 shape_len = 0
493 try:
494 shape_len = len(self.shape)
495 except TypeError:
496 pass
497
Louis Verhaard0411edb2020-11-16 16:37:11 +0100498 if shape_len > 4:
499 return
Tim Hall79d07d22020-04-27 18:20:16 +0100500 self.storage_rounding_quantum = arch.storage_rounding_quantums[self.format]
Louis Verhaard93719a92020-12-08 10:02:31 +0100501 self.storage_rounding_quantum = tuple(self.storage_rounding_quantum[-shape_len:])
Tim Hall79d07d22020-04-27 18:20:16 +0100502 self.brick_size = arch.brick_sizes[self.format]
Louis Verhaard93719a92020-12-08 10:02:31 +0100503 self.brick_size = tuple(self.brick_size[-shape_len:])
Tim Hall79d07d22020-04-27 18:20:16 +0100504 if self.shape is None:
505 return
506
507 self.bandwidth_shape = shape_round_to_quantum(self.shape, self.brick_size)
508 self.storage_shape = shape_round_to_quantum(self.shape, self.storage_rounding_quantum)
509
510 if fmt == TensorFormat.WeightsCompressed:
511 compression_ratio = 5 / 8
512 self.storage_compression_scale = compression_ratio
513 self.bandwidth_compression_scale = compression_ratio
514 self.compression_scale_for_worst_weight_stream = compression_ratio
515
Louis Verhaard93719a92020-12-08 10:02:31 +0100516 def storage_elements(self) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100517 elems = shape_num_elements(self.storage_shape)
518 if elems is None:
519 return 0
520 return elems
521
Louis Verhaard93719a92020-12-08 10:02:31 +0100522 def elements(self) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100523 elems = shape_num_elements(self.shape)
524 if elems is None:
525 return 0
526 return elems
527
Louis Verhaard93719a92020-12-08 10:02:31 +0100528 def has_fully_defined_shape(self) -> bool:
Tim Hall79d07d22020-04-27 18:20:16 +0100529 return shape_fully_defined(self.shape)
530
Louis Verhaard93719a92020-12-08 10:02:31 +0100531 def storage_size(self, scale: float = 1.0) -> int:
Patrik Gustavsson90831bc2020-08-24 16:26:11 +0200532 raw_size = self.storage_elements() * self.element_size() * scale
Tim Hall79d07d22020-04-27 18:20:16 +0100533 if raw_size == 0:
534 raw_size = 1 # force it to take up space
535 rounded_size = numeric_util.round_up(numeric_util.round_up_to_int(raw_size), self.alignment)
536 return rounded_size
537
Louis Verhaard93719a92020-12-08 10:02:31 +0100538 def storage_size_for_sub_purpose(
539 self, arch, sub_purpose: TensorSubPurpose, param_a: Optional[int] = None, param_b: Optional[int] = None
540 ) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100541 alt_shape = self.storage_shape_for_sub_purpose(sub_purpose, param_a, param_b)
542 elems = shape_num_elements(alt_shape)
543 if elems is None:
544 return 0
545 if sub_purpose == TensorSubPurpose.DoubleBuffer:
Patrik Gustavsson90831bc2020-08-24 16:26:11 +0200546 raw_size = (
547 elems
548 * self.element_size()
549 * self.compression_scale_for_worst_weight_stream
550 * arch.weight_estimation_scaling
551 )
Tim Hall79d07d22020-04-27 18:20:16 +0100552 else:
Patrik Gustavsson9baa4c32020-08-20 13:59:01 +0200553 # Rolling buffers are used for intermediate data in ifm streaming
554 # These will all use the NHCWB16 format, and need to be aligned to 16 in the C-dimension
555 if alt_shape[-1] % 16 != 0:
556 nhcwb16_shape = alt_shape[0:-1] + [numeric_util.round_up(alt_shape[-1], 16)]
557 elems = shape_num_elements(nhcwb16_shape)
558
Tim Hall79d07d22020-04-27 18:20:16 +0100559 raw_size = elems * self.element_size() * self.storage_compression_scale
560 rounded_size = numeric_util.round_up(numeric_util.round_up_to_int(raw_size), self.alignment)
561 return rounded_size
562
Louis Verhaard93719a92020-12-08 10:02:31 +0100563 def storage_shape_for_sub_purpose(
564 self, sub_purpose: TensorSubPurpose, param_a: Optional[int], param_b: Optional[int]
565 ) -> Shape:
Tim Hall79d07d22020-04-27 18:20:16 +0100566 if sub_purpose == TensorSubPurpose.DoubleBuffer:
Jacob Bohline843d332020-06-23 12:12:56 +0200567 shp = list(self.shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100568 assert len(shp) >= 2
Louis Verhaard93719a92020-12-08 10:02:31 +0100569 assert param_a is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100570 shp[-1] = min(shp[-1], param_a * 2)
Tim Hall79d07d22020-04-27 18:20:16 +0100571 else:
Jacob Bohline843d332020-06-23 12:12:56 +0200572 shp = list(self.storage_shape)
573 if sub_purpose == TensorSubPurpose.RollingBufferX:
574 assert len(shp) == 4
Louis Verhaard93719a92020-12-08 10:02:31 +0100575 assert param_a is not None
Jacob Bohline843d332020-06-23 12:12:56 +0200576 shp[0] = 1
577 shp[2] = min(shp[2], param_a)
578 elif sub_purpose == TensorSubPurpose.RollingBufferY:
579 assert len(shp) == 4
Louis Verhaard93719a92020-12-08 10:02:31 +0100580 assert param_a is not None
Jacob Bohline843d332020-06-23 12:12:56 +0200581 shp[0] = 1
582 shp[1] = min(shp[1], param_a)
583 elif sub_purpose == TensorSubPurpose.RollingBufferXY:
584 assert len(shp) == 4
Louis Verhaard93719a92020-12-08 10:02:31 +0100585 assert param_a is not None
586 assert param_b is not None
Jacob Bohline843d332020-06-23 12:12:56 +0200587 shp[0] = 1
588 shp[2] = min(shp[2], param_a)
589 shp[1] = min(shp[1], param_b)
590 elif sub_purpose == TensorSubPurpose.Standard:
591 pass
592 else:
593 assert 0, "did not expect new sub purpose %s" % (sub_purpose,)
594
Tim Hall79d07d22020-04-27 18:20:16 +0100595 return shp
596
Louis Verhaard93719a92020-12-08 10:02:31 +0100597 def set_new_sub_purpose(self, sub_purpose: TensorSubPurpose, param_a=None, param_b=None):
Tim Hall79d07d22020-04-27 18:20:16 +0100598 self.storage_shape = self.storage_shape_for_sub_purpose(sub_purpose, param_a, param_b)
599 self.sub_purpose = sub_purpose
600 if sub_purpose == TensorSubPurpose.DoubleBuffer:
601 self.storage_compression_scale = self.compression_scale_for_worst_weight_stream
602
Louis Verhaard93719a92020-12-08 10:02:31 +0100603 def bandwidth(self) -> float:
Tim Hall79d07d22020-04-27 18:20:16 +0100604 elems = shape_num_elements(self.bandwidth_shape)
605 if elems is None:
606 return 0
607 return elems * self.element_size() * self.bandwidth_compression_scale
608
Louis Verhaard93719a92020-12-08 10:02:31 +0100609 def consumers(self) -> List[Operation]:
Tim Hall79d07d22020-04-27 18:20:16 +0100610 return self.consumer_list
611
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000612 def addresses_for_rolling_buffer(self, start_coord: Shape, end_coord: Shape, fm_shape: Shape4D) -> Tuple:
Tim Hall79d07d22020-04-27 18:20:16 +0100613 # returns ( box_height0, box_height1, box_width, [address_tl, address_tr, address_bl, address_br] )
614
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100615 if self.storage_shape == []:
616 return (
617 1,
618 1,
619 1,
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000620 [self.address_for_coordinate(start_coord, shape=fm_shape.as_list()), None, None, None],
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100621 )
Tim Hall79d07d22020-04-27 18:20:16 +0100622
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100623 storage_shape_4D = full_shape(4, self.storage_shape, 1)
624 crossing_y = numeric_util.round_up(start_coord[1] + 1, storage_shape_4D[1])
625 crossing_x = numeric_util.round_up(start_coord[2] + 1, storage_shape_4D[2])
Tim Hall79d07d22020-04-27 18:20:16 +0100626
627 crossing_y = min(crossing_y, end_coord[1])
628 crossing_x = min(crossing_x, end_coord[2])
629
630 box_height0 = crossing_y - start_coord[1]
631 box_width = crossing_x - start_coord[2]
632
Louis Verhaard93719a92020-12-08 10:02:31 +0100633 addresses: List = [None] * 4
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000634 addresses[0] = self.address_for_coordinate(start_coord, shape=fm_shape.as_list())
Tim Hall79d07d22020-04-27 18:20:16 +0100635
636 if end_coord[2] > crossing_x:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100637 addresses[1] = self.address_for_coordinate(
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000638 [start_coord[0], start_coord[1], crossing_x, start_coord[3]], shape=fm_shape.as_list()
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100639 )
Michael McGeagh528a56d2020-12-16 11:33:21 +0000640 raise UnsupportedFeatureError("Striping in vertical direction is not supported")
Tim Hall79d07d22020-04-27 18:20:16 +0100641 if end_coord[1] > crossing_y:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100642 addresses[2] = self.address_for_coordinate(
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000643 [start_coord[0], crossing_y, start_coord[2], start_coord[3]], shape=fm_shape.as_list()
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100644 )
Tim Hall79d07d22020-04-27 18:20:16 +0100645 if end_coord[1] > crossing_y and end_coord[2] > crossing_x:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100646 addresses[3] = self.address_for_coordinate(
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000647 [start_coord[0], crossing_y, crossing_x, start_coord[3]], shape=fm_shape.as_list()
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100648 )
Tim Hall79d07d22020-04-27 18:20:16 +0100649
650 return box_height0, box_height0, box_width, addresses
651
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100652 def address_for_coordinate(self, coord: Shape, is_top_box: bool = False, shape: Shape = None) -> int:
653 if shape is None:
654 shape = self.shape
655 offset = self.address_offset_for_coordinate(coord, shape, is_top_box)
Louis Verhaard93719a92020-12-08 10:02:31 +0100656 assert offset is not None
657 return self.address + offset
Tim Hall79d07d22020-04-27 18:20:16 +0100658
Louis Verhaard93719a92020-12-08 10:02:31 +0100659 def get_strides_and_coord(self, coord: Optional[Shape] = None) -> Tuple[Optional[Shape], Optional[Shape]]:
Tim Hall79d07d22020-04-27 18:20:16 +0100660 if coord is None:
661 coord = [0] * len(self.storage_shape)
662
663 augmented_coord = coord
664 augmented_shape = self.storage_shape
665 while len(augmented_shape) < 4:
666 augmented_shape = [1] + augmented_shape
667
668 while len(augmented_coord) < 4:
669 augmented_coord = [0] + augmented_coord
670
671 assert len(augmented_coord) == len(augmented_shape)
672
673 if self.format == TensorFormat.NHWC:
674 augmented_shape = [augmented_shape[0], augmented_shape[3]] + augmented_shape[1:3] + [1]
675 augmented_coord = [augmented_coord[0], augmented_coord[3]] + augmented_coord[1:3] + [0]
Tim Hall79d07d22020-04-27 18:20:16 +0100676
677 elif self.format == TensorFormat.NHCWB16:
Patrik Gustavsson2213e902020-05-05 17:49:35 +0200678 channel_divisor = 16
Tim Hall79d07d22020-04-27 18:20:16 +0100679 augmented_shape = augmented_shape[0:4] + [1]
680 augmented_coord = (
681 [augmented_coord[0], augmented_coord[3] // channel_divisor]
682 + augmented_coord[1:3]
683 + [augmented_coord[3] % channel_divisor]
684 )
685
686 if augmented_shape[1] == 0:
687 augmented_shape[1] = 1
688
689 else:
Michael McGeaghf3e3ad72020-12-02 12:39:03 +0000690 assert self.format in (TensorFormat.Unknown, TensorFormat.WeightsCompressed)
Tim Hall79d07d22020-04-27 18:20:16 +0100691 return None, None
692
Louis Verhaard93719a92020-12-08 10:02:31 +0100693 strides: List = [0] * len(augmented_shape)
Tim Hall79d07d22020-04-27 18:20:16 +0100694 stride = self.element_size() * self.storage_compression_scale
695
696 if self.format != TensorFormat.NHCWB16:
Louis Verhaard93719a92020-12-08 10:02:31 +0100697 stride_order = [4, 1, 3, 2, 0]
Tim Hall79d07d22020-04-27 18:20:16 +0100698 for i in stride_order:
699 strides[i] = stride
700 stride *= augmented_shape[i]
701 else:
702 assert len(strides) == 5
Tim Hall79d07d22020-04-27 18:20:16 +0100703 strides[4] = stride
Patrik Gustavsson2213e902020-05-05 17:49:35 +0200704 strides[3] = 16 * stride # STRIDE_X
Tim Hall79d07d22020-04-27 18:20:16 +0100705 strides[1] = strides[3] * augmented_shape[2] # STRIDE_C
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200706 strides[2] = augmented_shape[2] * augmented_shape[3] * stride # STRIDE_Y
Tim Hall79d07d22020-04-27 18:20:16 +0100707 strides[0] = strides[2] * augmented_shape[1] # STRIDE_N
708
709 return strides, augmented_coord
710
Louis Verhaard93719a92020-12-08 10:02:31 +0100711 def get_strides(self) -> Shape:
Tim Hall79d07d22020-04-27 18:20:16 +0100712 strides, _ = self.get_strides_and_coord()
Louis Verhaard93719a92020-12-08 10:02:31 +0100713 assert strides is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100714 return strides
715
Louis Verhaard93719a92020-12-08 10:02:31 +0100716 def needs_dma(self) -> bool:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200717 return len(self.ops) == 1 and self.ops[0].type == Op.DMA
Louis Verhaard3c07c972020-05-07 08:12:58 +0200718
Louis Verhaard93719a92020-12-08 10:02:31 +0100719 def get_dma_src_tensor(self) -> "Optional[Tensor]":
Louis Verhaard3c07c972020-05-07 08:12:58 +0200720 # For weight tensors that need DMA: returns the source tensor in Flash, else None
721 # Note: for DMA ops, Pass.weight_tensor is referring to the SRAM weight tensor
722 return self.ops[0].inputs[0] if self.needs_dma() else None
723
Louis Verhaard93719a92020-12-08 10:02:31 +0100724 def find_npu_op(self) -> Optional[Operation]:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200725 # Returns the NPU operator that uses this tensor, excluding DMA operators.
726 for op in self.consumers():
Louis Verhaardaee5d752020-09-30 09:01:52 +0200727 if op.type == Op.DMA:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200728 return op.outputs[0].find_npu_op()
Dwight Lidman940fdee2020-08-13 13:11:48 +0200729 if op.run_on_npu:
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200730 return op
Louis Verhaard93719a92020-12-08 10:02:31 +0100731 return None
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200732
Louis Verhaard93719a92020-12-08 10:02:31 +0100733 def compressed_stream_index_from_coord(self, coord: Shape) -> int:
Tim Hall79d07d22020-04-27 18:20:16 +0100734 assert self.format == TensorFormat.WeightsCompressed
Louis Verhaard93719a92020-12-08 10:02:31 +0100735 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100736 assert len(self.compressed_values) > 0
737 assert len(self.compressed_values) + 1 == len(self.weight_compressed_offsets)
738
739 depth = coord[-1]
740 brick_depth = self.brick_size[-1]
741 # Clamp position at final element index
742 if depth > self.shape[-1]:
743 depth = self.shape[-1]
744
745 # Always round up to next boundary
Michael McGeagh8d3216f2020-08-10 11:35:57 +0100746 index = numeric_util.round_up_divide(depth, brick_depth)
Tim Hall79d07d22020-04-27 18:20:16 +0100747
748 # Check boundaries on all but last weight set (which may be shorter
749 # than the brick we divided it up into)
750 if index < len(self.weight_compressed_offsets) - 1:
751 # There are no half-way points in the weights
752 if (depth % brick_depth) != 0:
Michael McGeagh528a56d2020-12-16 11:33:21 +0000753 raise UnsupportedFeatureError("Offset into weights must be aligned to a brick")
Tim Hall79d07d22020-04-27 18:20:16 +0100754
755 return index
756
Louis Verhaard93719a92020-12-08 10:02:31 +0100757 def size_of_compressed_stream(self, index: int) -> int:
758 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100759 assert 0 <= index < len(self.compressed_values)
760 return len(self.compressed_values[index])
761
Louis Verhaard93719a92020-12-08 10:02:31 +0100762 def is_last_index_in_compressed_stream(self, index: int) -> bool:
763 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100764 assert 0 <= index < len(self.compressed_values)
765 return index == len(self.compressed_values) - 1
766
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100767 def address_offset_for_coordinate(self, orig_coord: Shape, shape: Shape, is_top_box: bool = False) -> Optional[int]:
Tim Hall79d07d22020-04-27 18:20:16 +0100768 address_offset = 0
769 coord = orig_coord
770
771 coord = coord[-len(self.storage_shape) :]
772
773 if self.sub_purpose == TensorSubPurpose.Standard:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100774 for idx, c in enumerate(orig_coord):
Tim Hall79d07d22020-04-27 18:20:16 +0100775 if is_top_box:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100776 assert c > 0 and c <= shape[idx]
Tim Hall79d07d22020-04-27 18:20:16 +0100777 else:
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100778 assert c >= 0 and c < shape[idx]
Tim Hall79d07d22020-04-27 18:20:16 +0100779
780 if self.format == TensorFormat.WeightsCompressed:
781 if len(self.weight_compressed_offsets) == 0:
782 return 0
783
Louis Verhaard3c07c972020-05-07 08:12:58 +0200784 if self.needs_dma() and self.sub_purpose == TensorSubPurpose.DoubleBuffer:
Tim Hall79d07d22020-04-27 18:20:16 +0100785 depth = orig_coord[-1]
786 brick_depth = self.brick_size[-1]
787 # Clamp position at final element index
788 if depth > self.shape[-1]:
789 depth = self.shape[-1]
790
791 # Always round up to next boundary
Michael McGeagh8d3216f2020-08-10 11:35:57 +0100792 index = numeric_util.round_up_divide(depth, brick_depth)
Tim Hall79d07d22020-04-27 18:20:16 +0100793 index = index % 2
Louis Verhaard93719a92020-12-08 10:02:31 +0100794 assert self.compressed_values is not None
Tim Hall79d07d22020-04-27 18:20:16 +0100795
796 if len(self.compressed_values) <= 2:
797 if is_top_box and index == 0:
798 for cv in self.compressed_values:
799 address_offset += len(cv)
800 else:
801 address_offset = index * len(self.compressed_values[0])
802 else:
803 if is_top_box and index == 0:
804 address_offset = self.storage_shape[-1]
805 else:
806 address_offset = index * (self.storage_shape[-1] // 2)
807 else:
808 index = self.compressed_stream_index_from_coord(orig_coord)
809 assert index < len(self.weight_compressed_offsets)
810 address_offset = self.weight_compressed_offsets[index]
811 else:
812 if is_top_box:
813 coord = [c - 1 for c in coord]
814
815 # handle wraparound for partial buffers. make sure to do this after subtracting top box:
816 coord = [c % self.storage_shape[idx] for idx, c in enumerate(coord)]
817
818 strides, augmented_coord = self.get_strides_and_coord(coord)
819 if strides is None:
820 return None
821
822 if is_top_box:
823 address_offset += 1 * strides[-1] # one element
824
825 address_offset += np.dot(augmented_coord, strides)
826
827 assert address_offset >= 0
828 assert address_offset <= self.storage_size()
829 return address_offset
830
Louis Verhaard93719a92020-12-08 10:02:31 +0100831 def is_allocated_in_tensor_arena(self, scratch_tensor_mem_area: MemArea) -> bool:
Michael McGeaghf3e3ad72020-12-02 12:39:03 +0000832 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 +0200833
Louis Verhaard93719a92020-12-08 10:02:31 +0100834 def equivalent(self, tens: "Tensor") -> bool:
Louis Verhaard0b8268a2020-08-05 16:11:29 +0200835 return self.equivalence_id == tens.equivalence_id
836
Louis Verhaard93719a92020-12-08 10:02:31 +0100837 def set_all_shapes(self, shape: Shape):
Michael McGeagh6a8d4242020-07-28 12:17:59 +0100838 self.shape = shape
839 self.storage_shape = shape
840 self.bandwidth_shape = shape
841
Louis Verhaard93719a92020-12-08 10:02:31 +0100842 def get_full_shape(self) -> Shape:
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100843 d = len(self.shape)
844 if d in (1, 3):
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100845 return full_shape(4, self.shape, 1)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100846 elif d == 2:
847 return [self.shape[0], 1, 1, self.shape[1]]
848 else:
Fredrik Svedberg835d8e12020-09-04 09:46:17 +0200849 return self.shape.copy()
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100850
Louis Verhaard93719a92020-12-08 10:02:31 +0100851 def is_quantized(self) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100852 # a tensor is quantized if it has an integral type and it contains valid quantization params
853
Tim Hall89567612020-10-27 11:57:57 +0000854 if not isinstance(self.quantization, QuantizationParameters):
Tim Hall93582962020-09-09 21:58:15 +0100855 return False
856
Tim Hall89567612020-10-27 11:57:57 +0000857 return (self.dtype.type & BaseType.Int) != 0 and self.quantization.is_valid()
Tim Hall93582962020-09-09 21:58:15 +0100858
Louis Verhaard6c74c3b2020-12-17 13:54:09 +0100859 def __lt__(self, other: "Tensor") -> bool:
860 return self.equivalence_id < other.equivalence_id
861
Tim Hall79d07d22020-04-27 18:20:16 +0100862 def __str__(self):
863 return "<nng.Tensor '%s' shape=%s dtype=%s>" % (self.name, self.shape, self.dtype)
864
865 __repr__ = __str__
Tim Hall93582962020-09-09 21:58:15 +0100866
Michael McGeagh528a56d2020-12-16 11:33:21 +0000867 def error(self, msg):
868 """
869 Raises a VelaError exception for errors encountered when parsing a Tensor
870
871 :param self: Tensor object that resulted in the error
872 :param msg: str object that contains a description of the specific error encountered
873 """
874
875 def _print_operators(ops):
876 lines = []
877 for idx, op in enumerate(ops):
878 op_type = getattr(op, "type", "Not an Operation")
879 op_id = getattr(op, "op_index", "-")
880 lines.append(f" {idx} = {op_type} ({op_id})")
881 return lines
882
883 lines = [f"Invalid {self.name} tensor. {msg}"]
884
885 lines += [" Driving operators:"]
886 lines += _print_operators(self.ops)
887
888 lines += [" Consuming operators:"]
889 lines += _print_operators(self.consumer_list)
890
891 raise VelaError("\n".join(lines))
892
Tim Hall93582962020-09-09 21:58:15 +0100893
Louis Verhaard93719a92020-12-08 10:02:31 +0100894def check_quantized_tens_scaling_equal(tens_a: Tensor, tens_b: Tensor) -> bool:
Tim Hall93582962020-09-09 21:58:15 +0100895 # checks that the scaling of two quantized tensors are equal
896
Tim Hall89567612020-10-27 11:57:57 +0000897 return tens_a.is_quantized() and tens_b.is_quantized() and tens_a.quantization.is_scaling_equal(tens_b.quantization)