blob: 312ac92e43b3d3a5ea2da47af98dde2d96a6fb06 [file] [log] [blame]
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +02001# Copyright (C) 2021 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.
16# Description:
17# TOSA mapping functions used by reader.
18# Contains a mapping from the various TOSA enums and options structs, generated by the FlatBuffer code
19# generator, to Vela's internal format.
20from .data_type import DataType
21from .operation import Op
22from .operation import TensorIndices
23from .tosa import ArithmeticRightShiftAttribute # noqa: F401
24from .tosa import AxisAttribute # noqa: F401
25from .tosa import ClampAttribute # noqa: F401
26from .tosa import CondIfAttribute # noqa: F401
27from .tosa import Conv2dAttribute # noqa: F401
28from .tosa import ConvQuantInfo # noqa: F401
29from .tosa import MatMulQuantInfo # noqa: F401
30from .tosa import MulAttribute # noqa: F401
31from .tosa import PadQuantInfo # noqa: F401
32from .tosa import Pool2dAttribute # noqa: F401
33from .tosa import ReluNAttribute # noqa: F401
34from .tosa import RescaleAttribute # noqa: F401
35from .tosa import ReshapeAttribute # noqa: F401
36from .tosa import ResizeAttribute # noqa: F401
37from .tosa import SliceAttribute # noqa: F401
38from .tosa import TileAttribute # noqa: F401
39from .tosa import TransposeConv2dAttribute # noqa: F401
40from .tosa import UnaryQuantInfo # noqa: F401
41from .tosa import WhileLoopAttribute # noqa: F401
42from .tosa.DType import DType
43from .tosa.Op import Op as TosaOp
44
45
46datatype_map = {
47 DType.BOOL: DataType.bool,
48 DType.UINT8: DataType.uint8,
49 DType.INT4: DataType.int4,
50 DType.INT8: DataType.int8,
51 DType.INT16: DataType.int16,
52 DType.INT32: DataType.int32,
53 DType.INT48: DataType.int48,
54 DType.FLOAT: DataType.float32,
55}
56
57
58# TODO duplicate of tflite_mapping
59def underscore_to_camel_case(s):
60 return "".join(x.title() for x in s.split("_"))
61
62
63# TODO duplicate of tflite_mapping
64def identity(x):
65 return x
66
67
68class AttrSerializer:
69 def __init__(self, name, members=None):
70 self.name = name
71 self.module = globals()[self.name]
72 self.cls = getattr(self.module, self.name)
73 self.members = []
74 if members is not None:
75 for mem in members:
76 deserialize = identity
77 is_vector = False
78 if isinstance(mem, tuple):
79 if len(mem) == 2:
80 mem, is_vector = mem
81 deserialize = tuple
82 else:
83 assert 0
84 underscore_mem = mem
85 camelcase_mem = underscore_to_camel_case(mem)
86 self.members.append((underscore_mem, camelcase_mem, deserialize, is_vector))
87
88 def deserialize(self, op_data):
89 attr_type = op_data.AttributeType()
90 attr = op_data.Attribute()
91 attrs = {}
92 if attr_type:
93 tosa_attrs = self.cls()
94 tosa_attrs.Init(attr.Bytes, attr.Pos)
95 for underscore_mem, camelcase_mem, deserialize, is_vector in self.members:
96 fun = camelcase_mem
97 if is_vector:
98 fun += "AsNumpy"
99
100 attr = getattr(tosa_attrs, fun)()
101 try:
102 attrs[underscore_mem] = deserialize(attr)
103 except TypeError:
104 print("Warning: {0} could not read attribute '{1}'.".format(self.name, underscore_mem))
105
106 return attrs
107
108
109class QuantSerializer:
110 def __init__(self, name, members=None):
111 self.name = name
112 self.module = globals()[self.name]
113 self.cls = getattr(self.module, self.name)
114 self.members = []
115 if members is not None:
116 for mem in members:
117 deserialize = identity
118 underscore_mem = mem
119 camelcase_mem = underscore_to_camel_case(mem)
120 self.members.append((underscore_mem, camelcase_mem, deserialize))
121
122 def deserialize(self, op_data):
123 quant_info_type = op_data.QuantInfoType()
124 quant_info = op_data.QuantInfo()
125 quant = {}
126 if quant_info_type:
127 tosa_quant = self.cls()
128 tosa_quant.Init(quant_info.Bytes, quant_info.Pos)
129 for underscore_mem, camelcase_mem, deserialize in self.members:
130 attr = getattr(tosa_quant, camelcase_mem)()
131 try:
132 quant[underscore_mem] = deserialize(attr)
133 except TypeError:
134 print("Warning: {0} could not read quant info '{1}'.".format(self.name, underscore_mem))
135
136 return quant
137
138
139is_vec = True
140pool2d_attrs = AttrSerializer("Pool2dAttribute", (("padding", is_vec), ("kernel", is_vec), ("stride", is_vec)))
141conv2d_attrs = AttrSerializer("Conv2dAttribute", (("padding", is_vec), ("stride", is_vec), ("dilation", is_vec)))
142transpose_conv2d_attrs = AttrSerializer(
143 "TransposeConv2dAttribute", (("outpad", is_vec), ("stride", is_vec), ("dilation", is_vec), ("out_shape", is_vec))
144)
145relun_attrs = AttrSerializer("ReluNAttribute", ("max_int"))
146axis_attrs = AttrSerializer("AxisAttribute", ("axis"))
147reshape_attrs = AttrSerializer("ReshapeAttribute", (("shape", is_vec),))
148slice_attrs = AttrSerializer("SliceAttribute", (("begin", is_vec), ("size", is_vec)))
149tile_attrs = AttrSerializer("TileAttribute", (("multiplies", is_vec),))
150resize_attrs = AttrSerializer(
151 "ResizeAttribute", (("output_size", is_vec), ("stride", is_vec), ("offset", is_vec), ("shift"))
152)
153clamp_attrs = AttrSerializer("ClampAttribute", (("min_int"), ("max_int")))
154rescale_attrs = AttrSerializer(
155 "RescaleAttribute",
156 ("input_zp", "output_zp", ("multiplier", is_vec), ("shift", is_vec), "scale32", "double_round", "per_channel"),
157)
158mul_attrs = AttrSerializer("MulAttribute", ("shift"))
159ars_attrs = AttrSerializer("ArithmeticRightShiftAttribute", ("round",))
160condif_attrs = AttrSerializer("CondIfAttribute", (("then_branch"), ("else_branch"))) # TODO these are references
161while_attrs = AttrSerializer("WhileLoopAttribute", (("cond_branch"), ("body_branch"))) # TODO these are references
162
163unary_quant_info = QuantSerializer("UnaryQuantInfo", ("input_zp", "output_zp"))
164conv_quant_info = QuantSerializer("ConvQuantInfo", ("input_zp", "weight_zp"))
165matmul_quant_info = QuantSerializer("MatMulQuantInfo", ("a_zp", "b_zp"))
166pad_quant_info = QuantSerializer("PadQuantInfo", ("input_zp"))
167
168unsupported_tosa_operators = {
169 TosaOp.UNKNOWN,
170 TosaOp.ARGMAX,
171 TosaOp.CONV3D,
172 TosaOp.MATMUL,
173 TosaOp.TRANSPOSE_CONV2D,
174 TosaOp.SIGMOID,
175 TosaOp.TANH,
176 TosaOp.BITWISE_AND,
177 TosaOp.BITWISE_OR,
178 TosaOp.BITWISE_XOR,
179 TosaOp.DIV,
180 TosaOp.LOGICAL_AND,
181 TosaOp.LOGICAL_LEFT_SHIFT,
182 TosaOp.LOGICAL_RIGHT_SHIFT,
183 TosaOp.LOGICAL_OR,
184 TosaOp.LOGICAL_XOR,
185 TosaOp.MAXIMUM,
186 TosaOp.MINIMUM,
187 TosaOp.MUL,
188 TosaOp.POW,
189 TosaOp.TABLE,
190 TosaOp.ABS,
191 TosaOp.BITWISE_NOT,
192 TosaOp.CEIL,
193 TosaOp.CLZ,
194 TosaOp.EXP,
195 TosaOp.FLOOR,
196 TosaOp.LOG,
197 TosaOp.LOGICAL_NOT,
198 TosaOp.NEGATE,
199 TosaOp.RECIPROCAL,
200 TosaOp.RSQRT,
201 TosaOp.SELECT,
202 TosaOp.EQUAL,
203 TosaOp.GREATER,
204 TosaOp.GREATER_EQUAL,
205 TosaOp.REDUCE_ANY,
206 TosaOp.REDUCE_ALL,
207 TosaOp.REDUCE_MAX,
208 TosaOp.REDUCE_MIN,
209 TosaOp.REDUCE_PRODUCT,
210 TosaOp.REDUCE_SUM,
211 TosaOp.CONCAT,
212 TosaOp.PAD,
213 TosaOp.RESHAPE,
214 TosaOp.REVERSE,
215 TosaOp.SLICE,
216 TosaOp.TILE,
217 TosaOp.TRANSPOSE,
218 TosaOp.GATHER,
219 TosaOp.SCATTER,
220 TosaOp.RESIZE,
221 TosaOp.CAST,
222 TosaOp.IDENTITY,
223 TosaOp.CUSTOM,
224 TosaOp.COND_IF,
225 TosaOp.WHILE_LOOP,
226}
227
228
229TOSA_NO_INDICES = TensorIndices([], [], [])
230TOSA_IFM_INDICES = TensorIndices([0], [], [])
231# TOSA_IFM_WEIGHTS_INDICES = TensorIndices([0], [1], [])
232TOSA_IFM_WEIGHTS_BIAS_INDICES = TensorIndices([0], [1], [2])
233TOSA_IFM_IFM2_INDICES = TensorIndices([0, 1], [], [])
234# TOSA_CONV2D_BACKPROP_INDICES = TensorIndices([2], [1], [3])
235# TOSA_TRANSPOSE_CONV_INDICES = TensorIndices([0], [1], [3])
236# TOSA_CONCAT_INDICES = TensorIndices([1, 2], [], [])
237# TOSA_SPLIT_IFM_INDICES = TensorIndices([1], [], [])
238# TOSA_BLOCK_LSTM_INDICES = TensorIndices([3], [4], [])
239
240
241tosa_operator_map = {
242 # TosaOp.UNKNOWN: (),
243 # TODO TosaOp.ARGMAX: (Op.ArgMax, axis_attrs, None),
244 TosaOp.AVG_POOL2D: (Op.AvgPool, pool2d_attrs, unary_quant_info, TOSA_IFM_INDICES),
245 TosaOp.CONV2D: (Op.Conv2DBias, conv2d_attrs, conv_quant_info, TOSA_IFM_WEIGHTS_BIAS_INDICES),
246 # TODO TosaOp.CONV3D:
247 TosaOp.DEPTHWISE_CONV2D: (Op.DepthwiseConv2DBias, conv2d_attrs, conv_quant_info, TOSA_IFM_WEIGHTS_BIAS_INDICES),
248 TosaOp.FULLY_CONNECTED: (Op.FullyConnected, None, conv_quant_info, TOSA_IFM_WEIGHTS_BIAS_INDICES),
249 # TODO TosaOp.MATMUL:
250 TosaOp.MAX_POOL2D: (Op.MaxPool, pool2d_attrs, None, TOSA_IFM_INDICES),
251 # TODO TosaOp.TRANSPOSE_CONV2D: (Op.Conv2DBackpropInput, transpose_conv2d_attrs, conv_quant_info)
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200252 TosaOp.CLAMP: (Op.Clamp, clamp_attrs, None, TOSA_IFM_INDICES),
Patrik Gustavsson8f1f9aa2021-06-28 07:41:58 +0200253 TosaOp.RELUN: (Op.ReluN, relun_attrs, None, TOSA_IFM_INDICES),
254 # TODO TosaOp.SIGMOID
255 # TODO TosaOp.TANH
256 TosaOp.ADD: (Op.Add, None, None, TOSA_IFM_IFM2_INDICES),
257 TosaOp.ARITHMETIC_RIGHT_SHIFT: (Op.SHR, ars_attrs, None, TOSA_IFM_IFM2_INDICES),
258 # TODO TosaOp.BITWISE_AND
259 # TODO TosaOp.BITWISE_OR
260 # TODO TosaOp.BITWISE_XOR
261 # TODO TosaOp.DIV
262 # TODO TosaOp.LOGICAL_AND
263 # TODO TosaOp.LOGICAL_LEFT_SHIFT
264 # TODO TosaOp.LOGICAL_RIGHT_SHIFT
265 # TODO TosaOp.LOGICAL_OR
266 # TODO TosaOp.LOGICAL_XOR
267 # TODO TosaOp.MAXIMUM
268 # TODO TosaOp.MINIMUM
269 # TODO TosaOp.MUL
270 # TODO TosaOp.POW
271 TosaOp.SUB: (Op.Sub, None, None, TOSA_IFM_IFM2_INDICES),
272 # TODO TosaOp.TABLE
273 # TODO TosaOp.ABS
274 # TODO TosaOp.BITWISE_NOT
275 # TODO TosaOp.CEIL
276 # TODO TosaOp.CLZ
277 # TODO TosaOp.EXP
278 # TODO TosaOp.FLOOR
279 # TODO TosaOp.LOG
280 # TODO TosaOp.LOGICAL_NOT
281 # TODO TosaOp.NEGATE
282 # TODO TosaOp.RECIPROCAL
283 # TODO TosaOp.RSQRT
284 # TODO TosaOp.SELECT
285 # TODO TosaOp.EQUAL
286 # TODO TosaOp.GREATER
287 # TODO TosaOp.GREATER_EQUAL
288 # TODO TosaOp.REDUCE_ANY
289 # TODO TosaOp.REDUCE_ALL
290 # TODO TosaOp.REDUCE_MAX
291 # TODO TosaOp.REDUCE_MIN
292 # TODO TosaOp.REDUCE_PRODUCT
293 # TODO TosaOp.REDUCE_SUM
294 # TODO TosaOp.CONCAT
295 # TODO TosaOp.PAD
296 # TODO TosaOp.RESHAPE
297 # TODO TosaOp.REVERSE
298 # TODO TosaOp.SLICE
299 # TODO TosaOp.TILE
300 # TODO TosaOp.TRANSPOSE
301 # TODO TosaOp.GATHER
302 # TODO TosaOp.SCATTER
303 # TODO TosaOp.RESIZE
304 # TODO TosaOp.CAST
305 TosaOp.RESCALE: (Op.Rescale, rescale_attrs, None, TOSA_IFM_INDICES),
306 TosaOp.CONST: (Op.Const, None, None, TOSA_NO_INDICES),
307 # TODO TosaOp.IDENTITY
308 # TODO TosaOp.CUSTOM
309 # TODO TosaOp.COND_IF
310 # TODO TosaOp.WHILE_LOOP
311}
312
313tosa_operator_inv_map = {v[0]: (k, v[1]) for k, v in tosa_operator_map.items()}
314
315
316def tosa_type_name(builtin):
317 return next(k for k, v in vars(TosaOp).items() if v == builtin)
318
319
320# TODO will return UNKNOWN for the once that have not yet been defined in tosa_operator_map
321def optype_to_tosa_op_type(op_type):
322 if op_type in tosa_operator_inv_map:
323 return tosa_type_name(tosa_operator_inv_map[op_type][0])
324 else:
325 return TosaOp.UNKNOWN