blob: f85cb4bbc6905150d7812d464dc85c1c491b7521 [file] [log] [blame]
Tim Halld0e41cf2023-02-14 14:54:18 +00001# SPDX-FileCopyrightText: Copyright 2020-2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
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.
Rickard Bolinbc6ee582022-11-04 08:24:29 +000016#
Tim Hall79d07d22020-04-27 18:20:16 +010017# Description:
18# Internal representation of a Neural Network Operation.
Jonas Ohlsson845e2322022-03-01 12:39:55 +010019# For Class name forward references for the type annotations. (see PEP 563).
20from __future__ import annotations
21
Louis Verhaarde8a5a782020-11-02 18:04:27 +010022import copy
Louis Verhaardaee5d752020-09-30 09:01:52 +020023from collections import namedtuple
24from enum import Enum
Dwight Lidman9b43f842020-12-08 17:56:44 +010025from typing import Any
26from typing import Dict
27from typing import List
Louis Verhaarde8a5a782020-11-02 18:04:27 +010028from typing import Optional
Louis Verhaardebf4af62021-01-27 15:57:57 +010029from typing import Tuple
Dwight Lidman9b43f842020-12-08 17:56:44 +010030from typing import TYPE_CHECKING
Tim Hall79d07d22020-04-27 18:20:16 +010031
Louis Verhaard1a92f782021-02-09 16:08:26 +010032from .api import NpuRoundingMode
Michael McGeagh528a56d2020-12-16 11:33:21 +000033from .errors import VelaError
Tim Hall3c5cfe92022-03-16 16:31:57 +000034from .ethos_u55_regs.ethos_u55_regs import resampling_mode
Tim Hall4ed38bc2020-10-20 18:54:20 +010035from .numeric_util import full_shape
patrik.gustavssoneeb85152020-12-21 17:10:40 +000036from .shape4d import Shape4D
Tim Hall4ed38bc2020-10-20 18:54:20 +010037
Jonas Ohlsson845e2322022-03-01 12:39:55 +010038# Import needed for Type annotations. Only import for Type checking to avoid run-time errors due to cyclic import.
Dwight Lidman9b43f842020-12-08 17:56:44 +010039if TYPE_CHECKING:
40 from .tensor import Tensor
41
Tim Hall4ed38bc2020-10-20 18:54:20 +010042PointXY = namedtuple("PointXY", "x y")
43PointXYZ = namedtuple("PointXYZ", "x y z")
44
Tim Hall79d07d22020-04-27 18:20:16 +010045
Louis Verhaardaee5d752020-09-30 09:01:52 +020046class NpuBlockType(Enum):
Tim Hall79d07d22020-04-27 18:20:16 +010047 Default = 0
48 ConvolutionMxN = 1
49 VectorProduct = 2
50 Pooling = 3
51 ConvolutionDepthWise = 4
52 ElementWise = 5
Fredrik Svedberga0c36242020-06-03 15:43:31 +020053 ReduceSum = 6
Tim Hall79d07d22020-04-27 18:20:16 +010054
55
Tim Hall4ed38bc2020-10-20 18:54:20 +010056class Kernel:
Louis Verhaarde8a5a782020-11-02 18:04:27 +010057 """
58 Kernel information for NPU operations
59 """
60
Tim Halld8339a72021-05-27 18:49:40 +010061 def __init__(
62 self,
63 w: int,
64 h: int,
65 stride_x: int = 1,
66 stride_y: int = 1,
67 dilation_x: int = 1,
68 dilation_y: int = 1,
69 valid_padding=False,
70 ):
Louis Verhaarde8a5a782020-11-02 18:04:27 +010071 assert stride_x > 0 and stride_y > 0
72 assert dilation_x > 0 and dilation_y > 0
Tim Hall4ed38bc2020-10-20 18:54:20 +010073 self.width = w
74 self.height = h
Louis Verhaarde8a5a782020-11-02 18:04:27 +010075 self.stride = PointXY(stride_x, stride_y)
76 self.dilation = PointXY(dilation_x, dilation_y)
Tim Halld8339a72021-05-27 18:49:40 +010077 self.valid_padding = valid_padding
Tim Hall4ed38bc2020-10-20 18:54:20 +010078
Louis Verhaarde8a5a782020-11-02 18:04:27 +010079 def elements_wh(self) -> int:
Tim Hall4ed38bc2020-10-20 18:54:20 +010080 return self.width * self.height
81
Louis Verhaarde8a5a782020-11-02 18:04:27 +010082 def area_width(self) -> int:
Tim Hall4ed38bc2020-10-20 18:54:20 +010083 return (self.width - 1) * self.dilation.x + 1
84
Louis Verhaarde8a5a782020-11-02 18:04:27 +010085 def area_height(self) -> int:
Tim Hall4ed38bc2020-10-20 18:54:20 +010086 return (self.height - 1) * self.dilation.y + 1
87
Louis Verhaardebf4af62021-01-27 15:57:57 +010088 def dilated_wh(self) -> Tuple[int, int]:
89 """Returns the dilated kernel width/height"""
90 return self.dilation.x * (self.width - 1) + 1, self.dilation.y * (self.height - 1) + 1
91
Louis Verhaarde8a5a782020-11-02 18:04:27 +010092 def __str__(self):
93 return f"w={self.width}, h={self.height}, stride={tuple(self.stride)}, dilation={tuple(self.dilation)}"
94
Tim Hall4ed38bc2020-10-20 18:54:20 +010095
Louis Verhaardaee5d752020-09-30 09:01:52 +020096# Classifies operators of type Custom
97class CustomType(Enum):
98 ThirdPartyOp = 0 # Third party custom op
99 NpuOp = 1 # NPU op
100 ExistingNpuOp = 2 # NPU op that was part of the input network
101
102
103TensorIndices = namedtuple("TensorIndices", ["ifms", "weights", "biases"])
104
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200105NNG_NO_INDICES = TensorIndices([], [], [])
106NNG_IFM_INDICES = TensorIndices([0], [], [])
107NNG_IFM_WEIGHTS_INDICES = TensorIndices([0], [1], [])
108NNG_IFM_WEIGHTS_BIAS_INDICES = TensorIndices([0], [1], [2])
109NNG_IFM_IFM2_INDICES = TensorIndices([0, 1], [], [])
110NNG_CONV2D_BACKPROP_INDICES = TensorIndices([2], [1], [3])
111NNG_TRANSPOSE_CONV_INDICES = TensorIndices([0], [1], [3])
112NNG_CONCAT_INDICES = TensorIndices([1, 2], [], [])
113NNG_SPLIT_IFM_INDICES = TensorIndices([1], [], [])
114NNG_BLOCK_LSTM_INDICES = TensorIndices([3], [4], [])
Louis Verhaardaee5d752020-09-30 09:01:52 +0200115
116
117# Static information related to operation codes
118class OperatorInfo:
119 __slots__ = ("id", "block_type", "indices", "is_unary")
120 _id = 0
121
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200122 def __init__(self, block_type=NpuBlockType.Default, indices=NNG_NO_INDICES, is_unary=False):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200123 OperatorInfo._id += 1
124 self.id = OperatorInfo._id
125 self.block_type = block_type
126 self.indices = indices # Indices of the different tensor purposes
127 self.is_unary = is_unary # Classifies elementwise operators
128
129
130# Internally used operation codes
131class Op(Enum):
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200132 Abs = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_INDICES, is_unary=True)
133 Add = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200134 AddN = OperatorInfo()
135 Any = OperatorInfo()
136 ArgMax = OperatorInfo()
137 ArgMin = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200138 AvgPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
erik.andersson@arm.com61f05d92022-09-27 12:06:32 +0200139 Atan2 = OperatorInfo(indices=NNG_IFM_IFM2_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200140 BatchMatMul = OperatorInfo()
141 BatchToSpaceND = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200142 BidirectionalSequenceLstm = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
143 BidirectionalSequenceRnn = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
144 BlockLSTM = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_BLOCK_LSTM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200145
146 CLZ = OperatorInfo(
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200147 block_type=NpuBlockType.ElementWise, indices=NNG_IFM_INDICES, is_unary=True
Louis Verhaardaee5d752020-09-30 09:01:52 +0200148 ) # NPU specific operation
149 Call = OperatorInfo()
150 Cast = OperatorInfo()
151 Ceil = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200152 Clamp = OperatorInfo(indices=NNG_IFM_INDICES) # TOSA specific
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100153 Clip = OperatorInfo() # NPU specific fused activation function for clipping between activation.min/max
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200154 Concat = OperatorInfo(indices=NNG_CONCAT_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200155 ConcatEmbeddings = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200156 ConcatSliceWrite = OperatorInfo(indices=NNG_IFM_INDICES)
157 ConcatTFLite = OperatorInfo(indices=NNG_CONCAT_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200158 Const = OperatorInfo() # Constant tensor, only used in CPU subgraphs
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200159 Conv2D = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=NNG_IFM_WEIGHTS_INDICES)
160 Conv2DBackpropInput = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=NNG_CONV2D_BACKPROP_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200161 Conv2DBackpropInputSwitchedBias = OperatorInfo(
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200162 block_type=NpuBlockType.ConvolutionMxN, indices=NNG_TRANSPOSE_CONV_INDICES
Louis Verhaardaee5d752020-09-30 09:01:52 +0200163 )
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200164 Conv2DBias = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=NNG_IFM_WEIGHTS_BIAS_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200165 Cos = OperatorInfo()
Tim Hall42abec12021-02-04 21:31:57 +0000166 Cumsum = OperatorInfo()
Louis Verhaardaee5d752020-09-30 09:01:52 +0200167 Custom = OperatorInfo() # Custom 3rd party operator, only used in CPU subgraphs
168 CustomNpuOp = OperatorInfo() # NPU custom operator, only used in CPU subgraphs
Louis Verhaardaee5d752020-09-30 09:01:52 +0200169 Delegate = OperatorInfo()
170 Densify = OperatorInfo()
171 DepthToSpace = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200172 DepthwiseConv2DBias = OperatorInfo(
173 block_type=NpuBlockType.ConvolutionDepthWise, indices=NNG_IFM_WEIGHTS_BIAS_INDICES
174 )
175 Dequantize = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200176 Div = OperatorInfo()
177 Elu = OperatorInfo()
178 EmbeddingLookup = OperatorInfo()
179 EmbeddingLookupSparse = OperatorInfo()
180 Equal = OperatorInfo()
181 Exp = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200182 ExpandDims = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200183 FakeQuantWithMinMaxArgs = OperatorInfo()
184 Fill = OperatorInfo()
185 Floor = OperatorInfo()
186 FloorDiv = OperatorInfo()
187 FloorMod = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200188 FullyConnected = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_BIAS_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200189 GatherNd = OperatorInfo()
190 GatherV2 = OperatorInfo()
191 Greater = OperatorInfo()
192 GreaterEqual = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200193 HardSwish = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200194 HashtableLookup = OperatorInfo()
Patrik Gustavssonef3ebdd2021-10-01 11:10:25 +0200195 Identity = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200196 If = OperatorInfo()
197 L2Norm = OperatorInfo()
198 L2Pool2D = OperatorInfo()
199 LRN = OperatorInfo()
200 LSHProjection = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200201 LeakyRelu = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_INDICES, is_unary=True)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200202 Less = OperatorInfo()
203 LessEqual = OperatorInfo()
204 Log = OperatorInfo()
205 LogSoftmax = OperatorInfo()
206 LogicalAnd = OperatorInfo()
207 LogicalNot = OperatorInfo()
208 LogicalOr = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200209 Lstm = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200210 LUT = OperatorInfo() # NPU specific, operator has LUT, only used in fused activation functions
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200211 MatMul = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200212 MatrixDiag = OperatorInfo()
213 MatrixSetDiag = OperatorInfo()
214 Max = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200215 MaxPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
216 Maximum = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES)
217 Mean = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200218 Min = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200219 Minimum = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200220 MirrorPad = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200221 Mul = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200222 Neg = OperatorInfo()
223 NonMaxSuppressionV4 = OperatorInfo()
224 NonMaxSuppressionV5 = OperatorInfo()
225 NotEqual = OperatorInfo()
226 OneHot = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200227 Pack = OperatorInfo(indices=NNG_IFM_INDICES)
228 PackReshaped = OperatorInfo(indices=NNG_IFM_INDICES)
229 Pad = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200230 PadV2 = OperatorInfo()
231 Placeholder = OperatorInfo() # Only used in CPU subgraphs
232 Pow = OperatorInfo()
Fredrik Svedberg8ddd4892022-08-19 16:06:04 +0200233 Prelu = OperatorInfo(indices=NNG_IFM_IFM2_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200234 Prod = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200235 Quantize = OperatorInfo(indices=NNG_IFM_INDICES)
236 QuantizedAvgPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
237 QuantizedConv2D = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=NNG_IFM_WEIGHTS_INDICES)
238 QuantizedMatMul = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
239 QuantizedMaxPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
240 QuantizedReshape = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200241 Range = OperatorInfo()
242 Rank = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200243 ReduceSum = OperatorInfo(block_type=NpuBlockType.ReduceSum, indices=NNG_IFM_INDICES)
244 Relu = OperatorInfo(indices=NNG_IFM_INDICES)
erik.andersson@arm.comdd49a722022-08-10 15:26:48 +0200245 Relu0To1 = OperatorInfo(indices=NNG_IFM_INDICES)
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200246 Relu6 = OperatorInfo(indices=NNG_IFM_INDICES)
247 ReluN1To1 = OperatorInfo(indices=NNG_IFM_INDICES)
248 ReluN = OperatorInfo(indices=NNG_IFM_INDICES) # TOSA specific
249 Rescale = OperatorInfo(indices=NNG_IFM_INDICES) # TOSA specific
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200250 Reshape = OperatorInfo(indices=NNG_IFM_INDICES)
Tim Hall885033b2022-07-21 11:46:03 +0100251 # resize ops map to pooling operations unless explicitly converted to other operations in the graph optimiser
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200252 ResizeBilinear = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
Tim Hall885033b2022-07-21 11:46:03 +0100253 ResizeNearestNeighbor = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200254 ReverseSequence = OperatorInfo()
255 ReverseV2 = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200256 Rnn = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200257 Round = OperatorInfo()
258 Rsqrt = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200259 SHL = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES) # NPU specific operation
260 SHR = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES) # NPU specific operation
Louis Verhaardaee5d752020-09-30 09:01:52 +0200261 ScatterNd = OperatorInfo()
262 SegmentSum = OperatorInfo()
263 Select = OperatorInfo()
264 SelectV2 = OperatorInfo()
Ayaan Masood4965fae2022-06-29 11:30:57 +0100265 Shape = OperatorInfo(indices=NNG_IFM_INDICES)
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200266 Sigmoid = OperatorInfo(indices=NNG_IFM_INDICES)
erik.andersson@arm.com61f05d92022-09-27 12:06:32 +0200267 Sign = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200268 SignBit = OperatorInfo()
269 Sin = OperatorInfo()
270 SkipGram = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200271 Slice = OperatorInfo(indices=NNG_IFM_INDICES)
272 Softmax = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200273 SpaceToBatchND = OperatorInfo()
274 SpaceToDepth = OperatorInfo()
275 SparseToDense = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200276 Split = OperatorInfo(indices=NNG_SPLIT_IFM_INDICES)
277 SplitSliceRead = OperatorInfo(indices=NNG_IFM_INDICES)
278 SplitV = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200279 Sqrt = OperatorInfo()
280 Square = OperatorInfo()
281 SquaredDifference = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200282 Squeeze = OperatorInfo(indices=NNG_IFM_INDICES)
283 StridedSlice = OperatorInfo(indices=NNG_IFM_INDICES)
284 Sub = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200285 SubgraphInput = OperatorInfo() # Only used in CPU subgraphs
286 Sum = OperatorInfo()
287 Svdf = OperatorInfo()
Patrik Gustavssonf436ada2021-09-14 14:56:48 +0200288 Table = OperatorInfo(indices=NNG_IFM_INDICES)
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200289 Tanh = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200290 Tile = OperatorInfo()
291 TopKV2 = OperatorInfo()
James Ward6bf16132021-09-08 11:14:20 +0100292 Transpose = OperatorInfo(indices=NNG_IFM_IFM2_INDICES)
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200293 UnidirectionalSequenceLstm = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
294 UnidirectionalSequenceRnn = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=NNG_IFM_WEIGHTS_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200295 Unique = OperatorInfo()
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200296 Unpack = OperatorInfo(indices=NNG_IFM_INDICES)
297 UnpackReshaped = OperatorInfo(indices=NNG_IFM_INDICES)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200298 Where = OperatorInfo()
299 While = OperatorInfo()
300 ZerosLike = OperatorInfo()
Dwight Lidman8a12da12021-07-19 13:43:05 +0200301 CallOnce = OperatorInfo()
302 BroadcastTo = OperatorInfo()
303 Rfft2D = OperatorInfo()
304 Conv3D = OperatorInfo()
305 Imag = OperatorInfo()
306 Real = OperatorInfo()
307 ComplexAbs = OperatorInfo()
308 Hashtable = OperatorInfo()
309 HashtableFind = OperatorInfo()
310 HashtableImport = OperatorInfo()
311 HashtableSize = OperatorInfo()
312 ReduceAll = OperatorInfo()
313 Conv3DTranspose = OperatorInfo()
Rickard Bolin2de898a2021-12-20 08:35:23 +0000314 VarHandle = OperatorInfo()
315 ReadVariable = OperatorInfo()
316 AssignVariable = OperatorInfo()
317 BroadcastArgs = OperatorInfo()
318 RandomStandardNormal = OperatorInfo()
Rickard Bolind66f8012022-04-21 07:36:55 +0000319 Bucketize = OperatorInfo()
320 RandomUniform = OperatorInfo()
321 Multinomial = OperatorInfo()
322 Gelu = OperatorInfo()
323 DynamicUpdateSlice = OperatorInfo()
erik.andersson@arm.comdd49a722022-08-10 15:26:48 +0200324 UnsortedSegmentProd = OperatorInfo()
erik.andersson@arm.com61f05d92022-09-27 12:06:32 +0200325 UnsortedSegmentMax = OperatorInfo()
326 UnsortedSegmentMin = OperatorInfo()
327 UnsortedSegmentSum = OperatorInfo()
Louis Verhaardaee5d752020-09-30 09:01:52 +0200328
329 @property
330 def info(self):
331 return self.value
332
333 @property
334 def npu_block_type(self):
335 return self.info.block_type
336
337 def is_conv2d_op(self):
338 return self.info.block_type == NpuBlockType.ConvolutionMxN
339
340 def is_depthwise_conv2d_op(self):
341 return self.info.block_type == NpuBlockType.ConvolutionDepthWise
342
343 def is_pool_op(self):
344 return self.info.block_type == NpuBlockType.Pooling
345
346 def is_maxpool_op(self):
347 return self in (Op.MaxPool, Op.QuantizedMaxPool)
348
349 def is_avgpool_op(self):
350 return self in (Op.QuantizedAvgPool, Op.AvgPool)
351
352 def is_elementwise_op(self):
353 return self.info.block_type == NpuBlockType.ElementWise
354
355 def is_unary_elementwise_op(self):
356 return self.info.block_type == NpuBlockType.ElementWise and self.info.is_unary
357
358 def is_binary_elementwise_op(self):
359 return self.info.block_type == NpuBlockType.ElementWise and not self.info.is_unary
360
361 def is_relu_op(self):
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200362 return self in (Op.Relu, Op.Relu6, Op.ReluN1To1, Op.ReluN, Op.Clip, Op.Clamp)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200363
364 def is_activation_op(self):
Diqing Zhong189f7482021-01-26 12:12:51 +0100365 return self.is_relu_op() or self in (Op.Tanh, Op.Sigmoid, Op.Softmax, Op.LUT, Op.HardSwish)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200366
367 def is_split_op(self):
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100368 return self in (Op.Split, Op.SplitV, Op.StridedSlice, Op.Slice, Op.UnpackReshaped, Op.Unpack)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200369
370 def is_concat_op(self):
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100371 return self in (Op.Concat, Op.ConcatTFLite, Op.PackReshaped, Op.Pack)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200372
Tim Hall885033b2022-07-21 11:46:03 +0100373 def is_resize_op(self):
374 return self in (Op.ResizeBilinear, Op.ResizeNearestNeighbor)
375
Louis Verhaardaee5d752020-09-30 09:01:52 +0200376 def needs_bias(self):
377 return bool(self.info.indices.biases)
378
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100379 def needs_shapes(self):
380 return bool(self.info.indices.ifms)
381
Louis Verhaardaee5d752020-09-30 09:01:52 +0200382 @classmethod
383 def op_set(cls, predicate):
384 # Returns the set of all operator codes that fulfill the given predicate
385 return {op_type for op_type in Op if predicate(op_type)}
386
387 def __str__(self):
388 return self.name
389
390 __repr__ = __str__
391
392 def __lt__(self, other):
393 return self.value.id < other.value.id
394
395
Michael McGeagh16895482020-12-14 15:51:20 +0000396class Padding(Enum):
397 SAME = 0
398 VALID = 1
Louis Verhaardae2d5532020-12-11 17:19:54 +0100399 EXPLICIT = 2 # Padding is specified in a PAD operation (only used for NPU operations)
Rickard Bolin9ae34552022-06-09 13:07:17 +0000400 TILE = 3 # Uses hardware tiles to pad by 1 with edge values on two sides of the IFM specified in explicit_padding
Michael McGeagh16895482020-12-14 15:51:20 +0000401
402
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100403class ActivationFunction:
404 """Fused activation function"""
405
406 def __init__(self, op_type: Op):
407 self.op_type = op_type # The activation operation to be performed
408 # min/max are optional; if present they are non-quantized values
409 self.min: Optional[float] = None
410 self.max: Optional[float] = None
411 # Table lookup index, only applicable for Op.LUT activation, 0-7
412 self.lut_index: int = 0
413
414 def clone(self):
415 res = copy.copy(self)
416 return res
417
418
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +0200419class ExplicitScaling:
420 """Explicit scaling parameters"""
421
422 def __init__(self, per_channel, shift, multiplier):
423 self.per_channel = per_channel
424 self.shift = shift
425 self.multiplier = multiplier
426
427 def clone(self):
428 res = copy.copy(self)
429 return res
430
431
432def create_activation_function(op_type: Op, min=None, max=None) -> ActivationFunction:
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100433 """Creates activation function with min/max depending on op_type"""
434 act = ActivationFunction(op_type)
435 if op_type == Op.Relu:
436 act.min = 0.0
437 elif op_type == Op.Relu6:
438 act.min = 0.0
439 act.max = 6.0
440 elif op_type == Op.ReluN1To1:
441 act.min = -1.0
442 act.max = 1.0
443 elif op_type == Op.Tanh:
444 act.min = -1.0
445 act.max = 1.0
446 elif op_type == Op.Sigmoid:
447 act.min = 0.0
448 act.max = 1.0
oliper01c4d35eb2022-06-21 08:51:01 +0000449 elif op_type == Op.Clamp:
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +0200450 assert min is not None and max is not None
451 act.min = min
452 act.max = max
453 elif op_type == Op.ReluN:
454 assert max is not None
455 act.min = 0.0
456 act.max = max
457
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100458 return act
459
460
Tim Hall79d07d22020-04-27 18:20:16 +0100461class Operation:
462 """Class representing a Neural Network operation. Has a name, a type,
Dwight Lidmanc6ac1942020-10-02 14:55:45 +0200463 input and output tensors, as well as an attribute dictionary."""
Tim Hall79d07d22020-04-27 18:20:16 +0100464
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200465 __slots__ = (
466 "type",
Rickard Bolinfea15162022-07-04 16:19:16 +0000467 "_original_type",
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200468 "name",
469 "op_index",
470 "attrs",
471 "inputs",
472 "outputs",
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100473 "intermediates",
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200474 "flops",
475 "scheduled_pass",
476 "run_on_npu",
Louis Verhaardaee5d752020-09-30 09:01:52 +0200477 "activation",
478 "memory_function",
Dwight Lidman4f728c02020-12-17 15:14:45 +0100479 "forced_input_quantization",
Louis Verhaardaee5d752020-09-30 09:01:52 +0200480 "forced_output_quantization",
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200481 "activation_lut",
Tim Hall4ed38bc2020-10-20 18:54:20 +0100482 "_kernel",
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100483 "ifm_shapes",
484 "ofm_shapes",
Fredrik Svedberge82be7c2021-01-18 15:21:03 +0100485 "rescale",
Patrik Gustavssone3b1b912021-02-09 15:38:46 +0100486 "read_offsets",
Tim Halld8339a72021-05-27 18:49:40 +0100487 "read_shapes",
Louis Verhaard1a92f782021-02-09 16:08:26 +0100488 "rounding_mode",
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +0200489 "explicit_scaling",
Dwight Lidman4f728c02020-12-17 15:14:45 +0100490 "low_precision_scaling",
Louis Verhaardc822d622021-03-11 14:59:06 +0100491 "write_offset",
492 "write_shape",
Tim Hall3c5cfe92022-03-16 16:31:57 +0000493 "ifm_resampling_mode",
Rickard Bolinfea15162022-07-04 16:19:16 +0000494 "tile_base_offsets_ifm",
495 "tile_base_offsets_ofm",
Rickard Bolin17e53b52022-09-06 16:09:01 +0000496 "ofm_stride_multiplier",
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200497 )
Tim Hall79d07d22020-04-27 18:20:16 +0100498
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100499 def __init__(self, op_type: Op, name: str):
Tim Hall79d07d22020-04-27 18:20:16 +0100500 self.type = op_type
Rickard Bolinfea15162022-07-04 16:19:16 +0000501 self._original_type = op_type # the original type of the operation. once set this shouldn't be changed
Tim Hall79d07d22020-04-27 18:20:16 +0100502 self.name = name
Dwight Lidman9b43f842020-12-08 17:56:44 +0100503 self.attrs: Dict[str, Any] = {}
Jonas Ohlsson845e2322022-03-01 12:39:55 +0100504 self.inputs: List[Optional[Tensor]] = []
Dwight Lidman9b43f842020-12-08 17:56:44 +0100505 self.outputs: List[Tensor] = []
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100506 self.intermediates: List[Tensor] = []
Tim Hall79d07d22020-04-27 18:20:16 +0100507 self.flops = 0
508 self.run_on_npu = True
Louis Verhaardaee5d752020-09-30 09:01:52 +0200509 # Fused activation function. If not none: operator code.
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100510 self.activation: Optional[ActivationFunction] = None
Louis Verhaardaee5d752020-09-30 09:01:52 +0200511 # Fused memory function, if not None: operator code
Louis Verhaardc822d622021-03-11 14:59:06 +0100512 self.memory_function: Optional[Op] = None
Louis Verhaardaee5d752020-09-30 09:01:52 +0200513 # If not none: contains QuantizationParameters to be used as output quantization
514 # (which overrides the ofm tensor's quantization), used in LUT
Dwight Lidman4f728c02020-12-17 15:14:45 +0100515 self.forced_input_quantization = None
Louis Verhaardaee5d752020-09-30 09:01:52 +0200516 self.forced_output_quantization = None
Tim Hall79d07d22020-04-27 18:20:16 +0100517 self.scheduled_pass = None
Tim Hallc8310b12020-06-17 14:53:11 +0100518 self.op_index = None # input network operator index
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200519 self.activation_lut = None
Tim Hall4ed38bc2020-10-20 18:54:20 +0100520 self._kernel = None
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000521 self.ifm_shapes: List[Shape4D] = []
522 self.ofm_shapes: List[Shape4D] = []
Jonas Ohlsson845e2322022-03-01 12:39:55 +0100523 self.read_offsets: List[Optional[Shape4D]] = [None, None] # offset for [ifm, ifm2]
524 self.read_shapes: List[Optional[Shape4D]] = [None, None] # read shape for [ifm, ifm2]
Louis Verhaard1a92f782021-02-09 16:08:26 +0100525 self.rounding_mode: Optional[NpuRoundingMode] = None
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +0200526 # Rescale op in TOSA supplies explicit multiplier and shift values
527 self.explicit_scaling: Optional[ExplicitScaling] = None
Dwight Lidman4f728c02020-12-17 15:14:45 +0100528 # The Mean operator (implemented as a depthwise convolution) requires scaling
529 # to be calculated differently in one case. In that case, this is set to True.
530 self.low_precision_scaling = False
Louis Verhaardc822d622021-03-11 14:59:06 +0100531 # Write offset, for operations that only produce a part of the OFM
532 self.write_offset: Optional[Shape4D] = None
533 # The amount of OFM that is produced by the operation (only if write_offset is not None).
534 # E.g. an operation that only fills the bottom row of an OFM of size 1x10x8x1 would have
535 # write_offset 0,9,0,0, write_shape 1,1,8,1
536 self.write_shape: Optional[Shape4D] = None
Tim Hall3c5cfe92022-03-16 16:31:57 +0000537 self.ifm_resampling_mode: resampling_mode = resampling_mode.NONE
Rickard Bolinfea15162022-07-04 16:19:16 +0000538 # ifm (nhwc), ifm2 (nhwc)
539 self.tile_base_offsets_ifm: List[List[int]] = [[0, 0, 0, 0], [0, 0, 0, 0]]
540 # ofm (nhwc)
541 self.tile_base_offsets_ofm: List[int] = [0, 0, 0, 0]
Rickard Bolin17e53b52022-09-06 16:09:01 +0000542 # For interleaved/sparse outputs - stride is multiplied with the stride factor of the corresponding axis
543 # Order is [C, H, W] - default is no multiplication
544 self.ofm_stride_multiplier: List[int] = [1, 1, 1]
Tim Hall79d07d22020-04-27 18:20:16 +0100545
546 def clone(self, suffix="_clone"):
547 res = Operation(self.type, self.name + suffix)
548
Rickard Bolinfea15162022-07-04 16:19:16 +0000549 # maintain the original type, in cases where the type was changed to something different
550 res._original_type = self._original_type
551
Tim Hall79d07d22020-04-27 18:20:16 +0100552 res.attrs = dict(self.attrs)
553 res.inputs = list(self.inputs)
554 res.outputs = list(self.outputs)
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100555 res.intermediates = list(self.intermediates)
Tim Hall79d07d22020-04-27 18:20:16 +0100556 res.flops = self.flops
Louis Verhaardaee5d752020-09-30 09:01:52 +0200557 res.run_on_npu = self.run_on_npu
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100558 res.activation = None if self.activation is None else self.activation.clone()
Louis Verhaardaee5d752020-09-30 09:01:52 +0200559 res.memory_function = self.memory_function
Dwight Lidman4f728c02020-12-17 15:14:45 +0100560 res.forced_input_quantization = self.forced_input_quantization
Louis Verhaardaee5d752020-09-30 09:01:52 +0200561 res.forced_output_quantization = self.forced_output_quantization
Tim Hall79d07d22020-04-27 18:20:16 +0100562 res.scheduled_pass = self.scheduled_pass
Tim Hallc8310b12020-06-17 14:53:11 +0100563 res.op_index = None # not relevant as not part of input network
Patrik Gustavssone3b1b912021-02-09 15:38:46 +0100564 res.read_offsets = list(self.read_offsets)
Tim Halld8339a72021-05-27 18:49:40 +0100565 res.read_shapes = list(self.read_shapes)
Rickard Bolinfea15162022-07-04 16:19:16 +0000566 res.write_offset = Shape4D(*self.write_offset) if self.write_offset else None
567 res.write_shape = Shape4D(*self.write_shape) if self.write_shape else None
Louis Verhaard1a92f782021-02-09 16:08:26 +0100568 res.rounding_mode = self.rounding_mode
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +0200569 res.explicit_scaling = self.explicit_scaling
Dwight Lidman4f728c02020-12-17 15:14:45 +0100570 res.low_precision_scaling = self.low_precision_scaling
Rickard Bolin814d01f2022-04-19 11:48:46 +0000571 res.ifm_resampling_mode = self.ifm_resampling_mode
Rickard Bolinfea15162022-07-04 16:19:16 +0000572 res.tile_base_offsets_ifm = [_ifm.copy() for _ifm in self.tile_base_offsets_ifm]
573 res.tile_base_offsets_ofm = self.tile_base_offsets_ofm.copy()
Rickard Bolin17e53b52022-09-06 16:09:01 +0000574 res.ofm_stride_multiplier = self.ofm_stride_multiplier.copy()
Tim Hall79d07d22020-04-27 18:20:16 +0100575
576 return res
577
578 def __str__(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200579 return "<nng.Operation '{}' type={}>".format(self.name, self.type)
Tim Hall79d07d22020-04-27 18:20:16 +0100580
581 __repr__ = __str__
582
Rickard Bolinfea15162022-07-04 16:19:16 +0000583 @property
584 def original_type(self):
585 return self._original_type
586
Fredrik Svedbergf3c7d552022-11-04 09:48:49 +0100587 @property
588 def type_changed(self):
589 return self.type != self.original_type
590
Michael McGeagh65fd9982020-10-20 11:49:28 +0100591 def get_kernel_size(self):
Tim Hall4ed38bc2020-10-20 18:54:20 +0100592 weights = self.weights
593 if weights and self.type.npu_block_type in (NpuBlockType.ConvolutionDepthWise, NpuBlockType.ConvolutionMxN):
594 weight_shape = full_shape(4, weights.shape, 1)
Michael McGeagh65fd9982020-10-20 11:49:28 +0100595 h = weight_shape[-4]
596 w = weight_shape[-3]
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100597 elif self.type.npu_block_type in (NpuBlockType.Pooling, NpuBlockType.ReduceSum) and "ksize" in self.attrs:
598 h, w = self.attrs["ksize"][1:3]
Tim Hall4ed38bc2020-10-20 18:54:20 +0100599 else:
Michael McGeagh65fd9982020-10-20 11:49:28 +0100600 h = self.attrs.get("filter_height", 1)
601 w = self.attrs.get("filter_width", 1)
602 return w, h
603
604 def get_kernel_stride(self):
605 if "strides" in self.attrs:
606 _, h, w, _ = self.attrs["strides"]
607 else:
608 h = self.attrs.get("stride_h", 1)
609 w = self.attrs.get("stride_w", 1)
610 return w, h
611
612 def get_kernel_dilation(self):
613 if "dilation" in self.attrs:
614 _, h, w, _ = self.attrs["dilation"]
615 else:
616 h = self.attrs.get("dilation_h_factor", 1)
617 w = self.attrs.get("dilation_w_factor", 1)
618 return w, h
619
620 @property
621 def kernel(self):
622 k_w, k_h = self.get_kernel_size()
623 s_w, s_h = self.get_kernel_stride()
624 d_w, d_h = self.get_kernel_dilation()
625 self._kernel = Kernel(k_w, k_h, s_w, s_h, d_w, d_h)
Tim Hall4ed38bc2020-10-20 18:54:20 +0100626 return self._kernel
627
Tim Hall79d07d22020-04-27 18:20:16 +0100628 def get_ifm_ifm2_weights_ofm(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200629 return self.ifm, self.ifm2, self.weights, self.ofm
Tim Hall79d07d22020-04-27 18:20:16 +0100630
Patrik Gustavssone2bfa7e2021-09-08 15:04:11 +0200631 def get_ifm_ifm2_ofm(self):
632 return self.ifm, self.ifm2, self.ofm
633
Tim Hall79d07d22020-04-27 18:20:16 +0100634 def get_ifm_weights_biases_ofm(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200635 return self.ifm, self.weights, self.bias, self.ofm
Tim Hall79d07d22020-04-27 18:20:16 +0100636
Jacob Bohlin49d92122020-08-19 14:36:46 +0200637 def get_ifm_ifm2_weights_biases_ofm(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200638 return self.ifm, self.ifm2, self.weights, self.bias, self.ofm
Jacob Bohlin49d92122020-08-19 14:36:46 +0200639
Louis Verhaardaee5d752020-09-30 09:01:52 +0200640 def get_ifm_ofm(self):
641 return self.ifm, self.ofm
Jacob Bohlin49d92122020-08-19 14:36:46 +0200642
Louis Verhaardaee5d752020-09-30 09:01:52 +0200643 @property
644 def ifm(self):
645 # Gets the IFM tensor, or None if not applicable
646 return self.get_input(self.type.info.indices.ifms, 0)
Jacob Bohlin49d92122020-08-19 14:36:46 +0200647
Louis Verhaardaee5d752020-09-30 09:01:52 +0200648 @property
649 def ifm2(self):
650 # Gets the IFM2 tensor, or None if not applicable
651 return self.get_input(self.type.info.indices.ifms, 1)
Louis Verhaard98a34992020-09-01 10:39:04 +0200652
Louis Verhaardaee5d752020-09-30 09:01:52 +0200653 @property
654 def bias(self):
655 # Gets the bias tensor, or None if not applicable
656 return self.get_input(self.type.info.indices.biases, 0)
657
658 @property
659 def weights(self):
660 # Gets the weight tensor, or None if not applicable
661 return self.get_input(self.type.info.indices.weights, 0)
662
663 def get_ifm_tensors(self):
664 # Gets the IFM tensors, or empty list if not applicable
665 return self._index_list_to_tensors(self.type.info.indices.ifms)
666
667 def get_weight_tensors(self):
668 # Gets the weight tensors, or empty list if not applicable
669 return self._index_list_to_tensors(self.type.info.indices.weights)
670
671 def get_bias_tensors(self):
672 # Gets the bias tensors, or empty list if not applicable
673 return self._index_list_to_tensors(self.type.info.indices.biases)
674
675 def _index_list_to_tensors(self, index_list):
676 return [self.inputs[ix] for ix in index_list if ix < len(self.inputs)]
677
678 def get_input(self, index_list, ix):
679 if ix >= len(index_list):
680 return None
681 if index_list[ix] >= len(self.inputs):
682 return None
683 return self.inputs[index_list[ix]]
684
685 @property
686 def ofm(self):
687 # Gets the OFM tensor, or None if not applicable
688 return self.outputs[0] if self.outputs else None
Tim Hall79d07d22020-04-27 18:20:16 +0100689
690 def get_concat_inputs_axis(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200691 assert self.type.is_concat_op()
Tim Hall79d07d22020-04-27 18:20:16 +0100692
Louis Verhaardaee5d752020-09-30 09:01:52 +0200693 if self.type == Op.Concat:
Tim Hall79d07d22020-04-27 18:20:16 +0100694 axis_tensor = self.inputs[0]
695 inputs = self.inputs[1:]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200696 elif self.type == Op.ConcatTFLite:
Tim Hall79d07d22020-04-27 18:20:16 +0100697 inputs = self.inputs
698 axis = self.attrs["axis"]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200699 elif self.type == Op.PackReshaped:
Tim Hall79d07d22020-04-27 18:20:16 +0100700 # Requires fixup_pack_input to be called before this point
701 inputs = self.inputs
702 axis = self.attrs["axis"]
703 assert len(self.inputs) == self.attrs["values_count"]
704 else:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200705 assert len(axis_tensor.ops) == 1 and axis_tensor.ops[0].type == Op.Const
Tim Hall79d07d22020-04-27 18:20:16 +0100706 axis = int(axis_tensor.values)
707
708 return inputs, axis
709
Louis Verhaardb2fb2122020-06-04 15:51:24 +0200710 def get_dilation_h_w(self):
711 _, dilation_h, dilation_w, _ = self.attrs.get("dilation", (1, 1, 1, 1))
712 return dilation_h, dilation_w
713
Tim Hall79d07d22020-04-27 18:20:16 +0100714 def get_split_inputs_axis(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200715 assert self.type.is_split_op()
Tim Hall79d07d22020-04-27 18:20:16 +0100716
717 offset_start = None
718 offset_end = None
719 axis = None
Louis Verhaardaee5d752020-09-30 09:01:52 +0200720 if self.type == Op.Split:
Tim Hall79d07d22020-04-27 18:20:16 +0100721 num_splits = self.attrs.get("num_splits")
722 axis_tens = self.inputs[0]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200723 assert len(axis_tens.ops) == 1 and axis_tens.ops[0].type == Op.Const
Tim Hall79d07d22020-04-27 18:20:16 +0100724 axis = int(axis_tens.values)
725 input_tens = self.inputs[1]
726 outputs = self.outputs
727 assert num_splits == len(outputs)
728
Louis Verhaardaee5d752020-09-30 09:01:52 +0200729 elif self.type == Op.SplitV:
Charles Xu53d47522020-05-04 11:32:05 +0200730 num_splits = self.attrs.get("num_splits")
731 input_tens = self.inputs[0]
732 size_tens = self.inputs[1]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200733 assert len(size_tens.ops) == 1 and size_tens.ops[0].type == Op.Const
Charles Xu53d47522020-05-04 11:32:05 +0200734 sizes = size_tens.values
Patrik Gustavsson271ddc32020-09-01 09:15:27 +0200735
Charles Xu53d47522020-05-04 11:32:05 +0200736 axis_tens = self.inputs[2]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200737 assert len(axis_tens.ops) == 1 and axis_tens.ops[0].type == Op.Const
Charles Xu53d47522020-05-04 11:32:05 +0200738 axis = int(axis_tens.values)
Patrik Gustavsson271ddc32020-09-01 09:15:27 +0200739
740 for idx, size in enumerate(sizes):
741 # One but only one size might be set to -1, indicating that size should be inferred
742 if size == -1:
743 sizes[idx] = input_tens.shape[axis] - (sum(sizes) + 1)
744 break
745
Charles Xu53d47522020-05-04 11:32:05 +0200746 outputs = self.outputs
747 assert num_splits == len(outputs)
748 assert sum(sizes) == input_tens.shape[axis]
749
Louis Verhaardaee5d752020-09-30 09:01:52 +0200750 elif self.type == Op.Slice:
Tim Hall79d07d22020-04-27 18:20:16 +0100751 input_tens, begin_tens, size_tens = self.inputs
752 outputs = self.outputs
753 offset_start = [0] * len(input_tens.shape)
754 offset_end = [0] * len(input_tens.shape)
755
756 for idx in range(len(begin_tens.values)):
Johan Alfvén0b799e42022-10-25 16:22:58 +0200757 offset_start[idx] = begin_tens.values[idx]
758 offset_end[idx] = size_tens.values[idx] + offset_start[idx]
Tim Hall79d07d22020-04-27 18:20:16 +0100759
Louis Verhaardaee5d752020-09-30 09:01:52 +0200760 elif self.type == Op.StridedSlice:
Tim Hall79d07d22020-04-27 18:20:16 +0100761 input_tens, begin_tens, end_tens, strides_tens = self.inputs
762 outputs = self.outputs
Tim Hall79d07d22020-04-27 18:20:16 +0100763
764 # Extract masks
Tim Hall79d07d22020-04-27 18:20:16 +0100765 ellipsis_mask = self.attrs["ellipsis_mask"]
Tim Hall79d07d22020-04-27 18:20:16 +0100766 new_axis_mask = self.attrs["new_axis_mask"]
767 shrink_axis_mask = self.attrs["shrink_axis_mask"]
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200768
769 # shrink_axis_mask/new_axis_mask/ellipsis_mask is not supported by the Operation class but the operation
Tim Hall79d07d22020-04-27 18:20:16 +0100770 # may have the attribute modified and handled in the graph optimization phase.
Patrik Gustavssoncf728902020-04-30 08:57:23 +0200771 assert shrink_axis_mask == new_axis_mask == ellipsis_mask == 0
Tim Halld0e41cf2023-02-14 14:54:18 +0000772 # use the begin and end values that were calculated in the model semantic check. this is because the end
773 # values can be affected (ignored) by the shrink_axis_mask and this mask may have been changed in the graph
774 # optimizer (see assert above)
775 offset_start = self.attrs["offset_begin"]
776 offset_end = self.attrs["offset_end"]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200777 elif self.type == Op.UnpackReshaped:
Tim Hall79d07d22020-04-27 18:20:16 +0100778 # Requires fixup_unpack_output to be called before this point
779 input_tens = self.inputs[0]
780 outputs = self.outputs
781 axis = self.attrs["axis"]
782 num_splits = self.attrs["num"]
783 # Number of outputs have to equal the value of the dimension to unpack
784 assert num_splits == len(outputs) == input_tens.shape[axis]
785 else:
786 assert False
787
788 return input_tens, outputs, axis, offset_start, offset_end
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200789
790 def set_activation_lut(self, lut_tensor):
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100791 self.activation = ActivationFunction(Op.LUT)
Fredrik Svedberga0c36242020-06-03 15:43:31 +0200792 self.activation_lut = lut_tensor
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100793 self.add_input_tensor(lut_tensor)
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100794
795 def add_input_tensor(self, tens):
796 self.inputs.append(tens)
797 if self not in tens.consumer_list:
798 tens.consumer_list.append(self)
799
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +0200800 def set_input_tensor(self, tens, idx):
801 tens_to_remove = self.inputs[idx]
802 if tens_to_remove in tens.consumer_list:
803 tens.consumer_list.remove(tens_to_remove)
804
805 self.inputs[idx] = tens
806 if self not in tens.consumer_list:
807 tens.consumer_list.append(self)
808
Dwight Lidman4f728c02020-12-17 15:14:45 +0100809 def get_input_quantization(self):
810 if self.forced_input_quantization is not None:
811 return self.forced_input_quantization
812 return self.ifm.quantization
813
Michael McGeagh5778ffd2020-08-06 17:31:02 +0100814 def set_output_tensor(self, tens):
815 tens.ops = [self]
816 self.outputs = [tens]
Jacob Bohlina41cd4d2020-08-26 18:21:28 +0200817
Louis Verhaard98a34992020-09-01 10:39:04 +0200818 def get_output_quantization(self):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200819 if self.forced_output_quantization is not None:
820 return self.forced_output_quantization
821 return self.ofm.quantization
Michael McGeagh528a56d2020-12-16 11:33:21 +0000822
823 def error(self, msg):
824 """
825 Raises a VelaError exception for errors encountered when parsing an Operation
826
827 :param self: Operation object that resulted in the error
828 :param msg: str object that contains a description of the specific error encountered
829 """
830
831 def _print_tensors(tensors):
832 lines = []
833 for idx, tens in enumerate(tensors):
834 tens_name = getattr(tens, "name", "Not a Tensor")
835 lines.append(f" {idx} = {tens_name}")
836 return lines
837
838 if self.op_index is None:
839 lines = [f"Invalid {self.type} (name = {self.name}) operator in the internal representation. {msg}"]
840 else:
841 lines = [f"Invalid {self.type} (op_index = {self.op_index}) operator in the input network. {msg}"]
842
843 lines += [" Input tensors:"]
844 lines += _print_tensors(self.inputs)
845
846 lines += [" Output tensors:"]
847 lines += _print_tensors(self.outputs)
848
849 raise VelaError("\n".join(lines))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100850
851 def set_ifm_ofm_shapes(self):
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000852 self.ifm_shapes = []
853 self.ofm_shapes = []
854
Fredrik Svedberg11563172022-07-06 14:54:12 +0200855 ifm_tensor, ifm2_tensor, ofm_tensor = self.get_ifm_ifm2_ofm()
856
857 if self.type == Op.Reshape:
858 # Set ofm shape
859 if len(self.inputs) > 1 and self.inputs[1].values is not None:
860 ofm_tensor.shape = self.inputs[1].values.flatten().tolist()
861 ofm_elements = ofm_tensor.elements()
862 # Stretch dimension
863 if ofm_elements < 0:
864 ofm_tensor.shape[ofm_tensor.shape.index(-1)] = int(ifm_tensor.elements() / abs(ofm_elements))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100865
866 # set all shapes to op, as 4D
867 if self.type == Op.FullyConnected:
Patrik Gustavsson2c2522d2021-01-29 11:51:31 +0100868 if len(self.ifm.shape) == 2:
869 self.ifm_shapes.append(Shape4D([self.ifm.shape[0], 1, 1, self.ifm.shape[1]]))
870 else:
871 # Special case, handled in graph optimization
872 self.ifm_shapes.append(Shape4D(ifm_tensor.get_full_shape()))
Johan Alfvén65835e02022-10-13 10:49:30 +0200873 self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape()))
874
Fredrik Svedberg11563172022-07-06 14:54:12 +0200875 elif self.type == Op.Softmax:
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000876 self.ifm_shapes.append(Shape4D(ifm_tensor.get_full_shape()))
877 self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape()))
Patrik Gustavssonda2b0032021-02-04 16:28:29 +0100878 elif self.type.is_split_op() or self.type.is_concat_op():
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100879 for inp in self.inputs:
880 if inp is not None:
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000881 self.ifm_shapes.append(Shape4D(full_shape(4, inp.shape, 1)))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100882 else:
883 self.ifm_shapes.append(None)
884 for out in self.outputs:
885 if out is not None:
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000886 self.ofm_shapes.append(Shape4D(full_shape(4, out.shape, 1)))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100887 else:
888 self.ofm_shapes.append(None)
889 else:
Patrik Gustavssonda2b0032021-02-04 16:28:29 +0100890 if ifm_tensor is not None:
891 self.ifm_shapes.append(Shape4D(full_shape(4, ifm_tensor.shape, 1)))
Patrik Gustavsson2349d422020-12-01 16:02:29 +0100892 if ifm2_tensor is not None:
patrik.gustavssoneeb85152020-12-21 17:10:40 +0000893 self.ifm_shapes.append(Shape4D(full_shape(4, ifm2_tensor.shape, 1)))
Patrik Gustavssonda2b0032021-02-04 16:28:29 +0100894 if ofm_tensor is not None:
895 self.ofm_shapes.append(Shape4D(full_shape(4, ofm_tensor.shape, 1)))
Tim Halld8339a72021-05-27 18:49:40 +0100896
897 def has_scaling(self):
898 scaled = True
899 for tensor in [self.ifm, self.ifm2, self.ofm]:
900 if tensor is not None:
901 if tensor.quantization is None:
902 scaled = False
903 break
904
905 return scaled