blob: 5ed9877790c1f632630eabee4532a99477cdd4f5 [file] [log] [blame]
Kevin Cheng3a478572021-01-22 17:21:02 -08001# Copyright (c) 2020-2021, ARM Limited.
Eric Kunzee5e26762020-10-13 16:11:07 -07002#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15#!/usr/bin/env python3
16
Kevin Cheng550ccc52021-03-03 11:21:43 -080017import os
18import sys
19import json
Eric Kunzee5e26762020-10-13 16:11:07 -070020import flatbuffers
21import numpy as np
22from enum import Enum, IntEnum, unique
Kevin Cheng550ccc52021-03-03 11:21:43 -080023from tosa import (
24 TosaGraph,
25 TosaBasicBlock,
26 TosaTensor,
27 TosaOperator,
28 DType,
29 Op,
30 ResizeMode,
31 Version,
32)
33
34# Include the ../thirdparty/serialization_lib/python directory in PYTHONPATH
35parent_dir = os.path.dirname(os.path.realpath(__file__))
36sys.path.append(
37 os.path.join(parent_dir, "..", "thirdparty", "serialization_lib", "python")
38)
Eric Kunzee5e26762020-10-13 16:11:07 -070039import tosa
Eric Kunzee5e26762020-10-13 16:11:07 -070040
41# With the way flatc generates its python types, there is no programatic way
42# to get string names for the integer types. Manually maintain a string table
43# here.
Kevin Cheng550ccc52021-03-03 11:21:43 -080044DTypeNames = [
45 "UNKNOWN",
46 "BOOL",
47 "UINT8",
48 "INT4",
49 "INT8",
50 "INT16",
51 "INT32",
52 "INT48",
53 "FLOAT",
54]
55
Eric Kunzee5e26762020-10-13 16:11:07 -070056
57def dtype_str_to_val(name):
58
59 for i in range(len(DTypeNames)):
60 if name.casefold() == DTypeNames[i].casefold():
61 return i
Kevin Cheng550ccc52021-03-03 11:21:43 -080062 raise Exception("Unable to parse DType name {}".format(name))
Eric Kunzee5e26762020-10-13 16:11:07 -070063
64
65class TosaSerializerUnion:
Kevin Cheng550ccc52021-03-03 11:21:43 -080066 """This class handles encapsulating and serializing union types into flatbuffers"""
67
Eric Kunzee5e26762020-10-13 16:11:07 -070068 def __init__(self):
69
70 # A tuple of the start and end functions. Set by the options constructors below
71 self.optFcns = None
72
73 # The type from the tosa.Options enumeration. Set by the options constructors below.
74 self.utype = None
75
76 # Each of these lists is a tuple of the add function and the
77 # value being added. Set by the options constructors below.
78 self.ints = []
79 self.bools = []
80 self.floats = []
81 self.strings = []
82 self.intvecs = []
Kevin Cheng77d0f762020-11-24 10:26:32 -080083 self.fpvecs = []
Eric Kunzee5e26762020-10-13 16:11:07 -070084
85 def serialize(self, builder):
86
87 # We have to build strings and vectors first
88 strList = []
89 intVecList = []
Kevin Cheng77d0f762020-11-24 10:26:32 -080090 fpVecList = []
Eric Kunzee5e26762020-10-13 16:11:07 -070091
92 for fcn, val in self.strings:
93 strList.append((fcn, builder.CreateString(val)))
94
95 for fcn, val in self.intvecs:
96 intVecList.append((fcn, TosaSerializer.serializeInt32Vec(builder, val)))
97
Kevin Cheng77d0f762020-11-24 10:26:32 -080098 for fcn, val in self.fpvecs:
99 fpVecList.append((fcn, TosaSerializer.serializeFpVec(builder, val)))
100
Eric Kunzee5e26762020-10-13 16:11:07 -0700101 startFcn, endFcn = self.optFcns
102
103 # Then serialize the options object from the list of primitives and
104 # other serialized values
105 startFcn(builder)
106 for fcn, val in self.ints:
107 fcn(builder, val)
108
109 for fcn, val in self.bools:
110 fcn(builder, val)
111
112 for fcn, val in self.floats:
113 fcn(builder, val)
114
115 for fcn, val in strList:
116 fcn(builder, val)
117
118 for fcn, val in intVecList:
119 fcn(builder, val)
120
Kevin Cheng77d0f762020-11-24 10:26:32 -0800121 for fcn, val in fpVecList:
122 fcn(builder, val)
123
Eric Kunzee5e26762020-10-13 16:11:07 -0700124 return endFcn(builder)
125
Kevin Cheng550ccc52021-03-03 11:21:43 -0800126
Eric Kunzee5e26762020-10-13 16:11:07 -0700127class TosaSerializerAttribute(TosaSerializerUnion):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800128 """This class handles encapsulating all of the enumerated types for attributes"""
Eric Kunzee5e26762020-10-13 16:11:07 -0700129
130 def __init__(self):
131 super().__init__()
132
133 def Pool2dAttribute(self, kernel, stride, padding):
134 from tosa import Pool2dAttribute as a, Attribute
135
136 self.utype = Attribute.Attribute().Pool2dAttribute
137
138 self.optFcns = (a.Pool2dAttributeStart, a.Pool2dAttributeEnd)
Kevin Cheng550ccc52021-03-03 11:21:43 -0800139 self.intvecs.append((a.Pool2dAttributeAddPadding, padding))
140 self.intvecs.append((a.Pool2dAttributeAddKernel, kernel))
141 self.intvecs.append((a.Pool2dAttributeAddStride, stride))
Eric Kunzee5e26762020-10-13 16:11:07 -0700142
143 def Conv2dAttribute(self, padding, stride, dilation):
144 from tosa import Conv2dAttribute as a, Attribute
145
146 self.utype = Attribute.Attribute().Conv2dAttribute
147 self.optFcns = (a.Conv2dAttributeStart, a.Conv2dAttributeEnd)
148
Kevin Cheng550ccc52021-03-03 11:21:43 -0800149 self.intvecs.append((a.Conv2dAttributeAddPadding, padding))
150 self.intvecs.append((a.Conv2dAttributeAddStride, stride))
151 self.intvecs.append((a.Conv2dAttributeAddDilation, dilation))
Eric Kunzee5e26762020-10-13 16:11:07 -0700152
153 def TransposeConv2DAttribute(self, outpad, stride, dilation, output_shape):
154 from tosa import TransposeConv2dAttribute as a, Attribute
155
156 self.utype = Attribute.Attribute().TransposeConv2dAttribute
157 self.optFcns = (a.TransposeConv2dAttributeStart, a.TransposeConv2dAttributeEnd)
158
Kevin Cheng550ccc52021-03-03 11:21:43 -0800159 self.intvecs.append((a.TransposeConv2dAttributeAddOutpad, outpad))
160 self.intvecs.append((a.TransposeConv2dAttributeAddStride, stride))
161 self.intvecs.append((a.TransposeConv2dAttributeAddDilation, dilation))
162 self.intvecs.append((a.TransposeConv2dAttributeAddOutputShape, output_shape))
Eric Kunzee5e26762020-10-13 16:11:07 -0700163
164 def ReluNAttribute(self, maxint, maxfp):
165 from tosa import ReluNAttribute as a, Attribute
166
167 self.utype = Attribute.Attribute().ReluNAttribute
168 self.optFcns = (a.ReluNAttributeStart, a.ReluNAttributeEnd)
169
170 self.ints.append((a.ReluNAttributeAddMaxInt, maxint))
171 self.ints.append((a.ReluNAttributeAddMaxFp, maxfp))
172
Eric Kunzee5e26762020-10-13 16:11:07 -0700173 def AxisAttribute(self, axis):
174 from tosa import AxisAttribute as a, Attribute
175
176 self.utype = Attribute.Attribute().AxisAttribute
177 self.optFcns = (a.AxisAttributeStart, a.AxisAttributeEnd)
178
Kevin Cheng550ccc52021-03-03 11:21:43 -0800179 self.ints.append((a.AxisAttributeAddAxis, axis))
Eric Kunzee5e26762020-10-13 16:11:07 -0700180
181 def ReshapeAttribute(self, shape):
182 from tosa import ReshapeAttribute as a, Attribute
183
184 self.utype = Attribute.Attribute().ReshapeAttribute
185 self.optFcns = (a.ReshapeAttributeStart, a.ReshapeAttributeEnd)
186
Kevin Cheng550ccc52021-03-03 11:21:43 -0800187 self.intvecs.append((a.ReshapeAttributeAddShape, shape))
Eric Kunzee5e26762020-10-13 16:11:07 -0700188
189 def SliceAttribute(self, begin, size):
190 from tosa import SliceAttribute as a, Attribute
191
192 self.utype = Attribute.Attribute().SliceAttribute
193 self.optFcns = (a.SliceAttributeStart, a.SliceAttributeEnd)
194
Kevin Cheng550ccc52021-03-03 11:21:43 -0800195 self.intvecs.append((a.SliceAttributeAddBegin, begin))
196 self.intvecs.append((a.SliceAttributeAddSize, size))
Eric Kunzee5e26762020-10-13 16:11:07 -0700197
198 def TileAttribute(self, multiples):
199 from tosa import TileAttribute as a, Attribute
200
201 self.utype = Attribute.Attribute().TileAttribute
202 self.optFcns = (a.TileAttributeStart, a.TileAttributeEnd)
203
Kevin Cheng550ccc52021-03-03 11:21:43 -0800204 self.intvecs.append((a.TileAttributeAddMultiples, multiples))
Eric Kunzee5e26762020-10-13 16:11:07 -0700205
Kevin Cheng550ccc52021-03-03 11:21:43 -0800206 def ResizeAttribute(
207 self, output_size, stride, offset, shift, stride_fp, offset_fp, mode
208 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700209 from tosa import ResizeAttribute as a, Attribute
210
211 self.utype = Attribute.Attribute().ResizeAttribute
212 self.optFcns = (a.ResizeAttributeStart, a.ResizeAttributeEnd)
213
Kevin Cheng550ccc52021-03-03 11:21:43 -0800214 self.intvecs.append((a.ResizeAttributeAddOutputSize, output_size))
215 self.intvecs.append((a.ResizeAttributeAddStride, stride))
216 self.intvecs.append((a.ResizeAttributeAddOffset, offset))
217 self.ints.append((a.ResizeAttributeAddShift, shift))
218 self.fpvecs.append((a.ResizeAttributeAddStrideFp, stride_fp))
219 self.fpvecs.append((a.ResizeAttributeAddOffsetFp, offset_fp))
220 self.ints.append((a.ResizeAttributeAddMode, mode))
Eric Kunzee5e26762020-10-13 16:11:07 -0700221
222 def ClampAttribute(self, minint, maxint, minfp, maxfp):
223 from tosa import ClampAttribute as a, Attribute
224
225 self.utype = Attribute.Attribute().ClampAttribute
226 self.optFcns = (a.ClampAttributeStart, a.ClampAttributeEnd)
227
Kevin Cheng550ccc52021-03-03 11:21:43 -0800228 self.ints.append((a.ClampAttributeAddMinInt, minint))
229 self.ints.append((a.ClampAttributeAddMaxInt, maxint))
Eric Kunzee5e26762020-10-13 16:11:07 -0700230
Kevin Cheng550ccc52021-03-03 11:21:43 -0800231 self.ints.append((a.ClampAttributeAddMinFp, minfp))
232 self.ints.append((a.ClampAttributeAddMaxFp, maxfp))
Eric Kunzee5e26762020-10-13 16:11:07 -0700233
Kevin Cheng550ccc52021-03-03 11:21:43 -0800234 def RescaleAttribute(
235 self, input_zp, output_zp, multiplier, shift, scale32, double_round, per_channel
236 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700237 from tosa import RescaleAttribute as a, Attribute
238
239 self.utype = Attribute.Attribute().RescaleAttribute
240 self.optFcns = (a.RescaleAttributeStart, a.RescaleAttributeEnd)
241
Kevin Cheng550ccc52021-03-03 11:21:43 -0800242 self.ints.append((a.RescaleAttributeAddInputZp, input_zp))
243 self.ints.append((a.RescaleAttributeAddOutputZp, output_zp))
244 self.intvecs.append((a.RescaleAttributeAddMultiplier, multiplier))
245 self.intvecs.append((a.RescaleAttributeAddShift, shift))
246 self.bools.append((a.RescaleAttributeAddScale32, scale32))
247 self.bools.append((a.RescaleAttributeAddDoubleRound, double_round))
248 self.bools.append((a.RescaleAttributeAddPerChannel, per_channel))
Eric Kunzee5e26762020-10-13 16:11:07 -0700249
Kevin Chengaee1fac2020-11-11 13:54:06 -0800250 def MulAttribute(self, shift):
251 from tosa import MulAttribute as a, Attribute
252
253 self.utype = Attribute.Attribute().MulAttribute
254 self.optFcns = (a.MulAttributeStart, a.MulAttributeEnd)
255
Kevin Cheng550ccc52021-03-03 11:21:43 -0800256 self.ints.append((a.MulAttributeAddShift, shift))
Kevin Chengaee1fac2020-11-11 13:54:06 -0800257
258 def ArithmeticRightShiftAttribute(self, round):
259 from tosa import ArithmeticRightShiftAttribute as a, Attribute
260
261 self.utype = Attribute.Attribute().ArithmeticRightShiftAttribute
Kevin Cheng550ccc52021-03-03 11:21:43 -0800262 self.optFcns = (
263 a.ArithmeticRightShiftAttributeStart,
264 a.ArithmeticRightShiftAttributeEnd,
265 )
Kevin Chengaee1fac2020-11-11 13:54:06 -0800266
Kevin Cheng550ccc52021-03-03 11:21:43 -0800267 self.bools.append((a.ArithmeticRightShiftAttributeAddRound, round))
Kevin Chengaee1fac2020-11-11 13:54:06 -0800268
Eric Kunzee5e26762020-10-13 16:11:07 -0700269 def CustomAttribute(self, identifier):
270 from tosa import CustomAttribute as a, Attribute
271
272 self.utype = Attribute.Attribute().CustomAttribute
273 self.optFcns = (a.CustomAttributeStart, a.CustomAttributeEnd)
274
Kevin Cheng550ccc52021-03-03 11:21:43 -0800275 self.strings.append((a.CustomAttributeAddIdentifier, identifier))
Eric Kunzee5e26762020-10-13 16:11:07 -0700276
277 def CondIfAttribute(self, then_branch, else_branch):
278 from tosa import CondIfAttribute as a, Attribute
279
280 self.utype = Attribute.Attribute().CondIfAttribute
281 self.optFcns = (a.CondIfAttributeStart, a.CondIfAttributeEnd)
282
Kevin Cheng550ccc52021-03-03 11:21:43 -0800283 self.strings.append((a.CondIfAttributeAddThenBranch, then_branch))
284 self.strings.append((a.CondIfAttributeAddElseBranch, else_branch))
Eric Kunzee5e26762020-10-13 16:11:07 -0700285
286 def WhileLoopAttribute(self, cond_branch, body_branch):
287 from tosa import WhileLoopAttribute as a, Attribute
288
289 self.utype = Attribute.Attribute().WhileLoopAttribute
290 self.optFcns = (a.WhileLoopAttributeStart, a.WhileLoopAttributeEnd)
291
Kevin Cheng550ccc52021-03-03 11:21:43 -0800292 self.strings.append((a.WhileLoopAttributeAddCondBranch, cond_branch))
293 self.strings.append((a.WhileLoopAttributeAddBodyBranch, body_branch))
294
Eric Kunzee5e26762020-10-13 16:11:07 -0700295
296class TosaSerializerQuantInfo(TosaSerializerUnion):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800297 """This class handles encapsulating all of the enumerated types for quantinfo types"""
298
Eric Kunzee5e26762020-10-13 16:11:07 -0700299 def __init__(self):
300 super().__init__()
301
302 def ConvQuantInfo(self, input_zp, weight_zp):
303 from tosa import ConvQuantInfo as q, QuantInfo
304
305 self.utype = QuantInfo.QuantInfo().ConvQuantInfo
306 self.optFcns = (q.ConvQuantInfoStart, q.ConvQuantInfoEnd)
307 self.ints.append((q.ConvQuantInfoAddInputZp, input_zp))
308 self.ints.append((q.ConvQuantInfoAddWeightZp, weight_zp))
309
310 def UnaryQuantInfo(self, input_zp, output_zp):
311 from tosa import UnaryQuantInfo as q, QuantInfo
312
313 self.utype = QuantInfo.QuantInfo().UnaryQuantInfo
314 self.optFcns = (q.UnaryQuantInfoStart, q.UnaryQuantInfoEnd)
315 self.ints.append((q.UnaryQuantInfoAddInputZp, input_zp))
316 self.ints.append((q.UnaryQuantInfoAddOutputZp, output_zp))
317
318 def MatMulQuantInfo(self, a_zp, b_zp):
319 from tosa import MatMulQuantInfo as q, QuantInfo
320
321 self.utype = QuantInfo.QuantInfo().MatMulQuantInfo
322 self.optFcns = (q.MatMulQuantInfoStart, q.MatMulQuantInfoEnd)
323 self.ints.append((q.MatMulQuantInfoAddAZp, a_zp))
324 self.ints.append((q.MatMulQuantInfoAddBZp, b_zp))
325
326 def PadQuantInfo(self, input_zp):
327 from tosa import PadQuantInfo as q, QuantInfo
328
329 self.utype = QuantInfo.QuantInfo().PadQuantInfo
330 self.optFcns = (q.PadQuantInfoStart, q.PadQuantInfoEnd)
331 self.ints.append((q.PadQuantInfoAddInputZp, input_zp))
332
Kevin Cheng550ccc52021-03-03 11:21:43 -0800333
Eric Kunzee5e26762020-10-13 16:11:07 -0700334class TosaSerializerTensor:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800335 def __init__(
336 self,
337 name,
338 shape,
339 dtype,
340 filename=None,
341 placeholderFilename=None,
342 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700343 self.name = name
344
345 if isinstance(shape, np.ndarray):
346 shape = shape.astype(int).tolist()
347 shape = list(map(int, shape))
348
349 self.shape = shape
350 self.dtype = dtype
Eric Kunzee5e26762020-10-13 16:11:07 -0700351
352 # Filename for const tensors. This gets written to the .tosa serialization
353 self.filename = filename
354
355 # Filename for placeholder tensors. These get generated by the test generation
356 # process and are written to disk, but are considered input tensors by the network
357 # so they do not appear in the TOSA serialiazation. However, if we want to form a unit
358 # test around these input tensors, we can get the filename from here.
359 self.placeholderFilename = placeholderFilename
360
361 def __str__(self):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800362 str = "TosaSerializerTensor name: {} shape: {} dtype: {} filename: {}".format(
363 self.name,
364 self.shape,
365 DTypeNames[self.dtype],
366 self.filename,
367 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700368 return str
369
Eric Kunzee5e26762020-10-13 16:11:07 -0700370 def setDtype(self, dtype):
371 self.dtype = dtype
372
Eric Kunzee5e26762020-10-13 16:11:07 -0700373 def serialize(self, builder):
374 fb_name = builder.CreateString(self.name)
375 if self.filename:
376 fb_filename = builder.CreateString(self.filename)
377 fb_shapes = TosaSerializer.serializeInt32Vec(builder, self.shape)
Eric Kunzee5e26762020-10-13 16:11:07 -0700378
379 TosaTensor.TosaTensorStart(builder)
380 TosaTensor.TosaTensorAddName(builder, fb_name)
381 TosaTensor.TosaTensorAddShape(builder, fb_shapes)
382 TosaTensor.TosaTensorAddType(builder, self.dtype)
Eric Kunzee5e26762020-10-13 16:11:07 -0700383 if self.filename:
384 TosaTensor.TosaTensorAddNpyFilename(builder, fb_filename)
385
386 return TosaTensor.TosaTensorEnd(builder)
387
Kevin Cheng550ccc52021-03-03 11:21:43 -0800388
Eric Kunzee5e26762020-10-13 16:11:07 -0700389class TosaSerializerOperator:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800390 def __init__(self, op, inputs, outputs, attributes=None, quantInfo=None):
Eric Kunzee5e26762020-10-13 16:11:07 -0700391 self.op = op
392 self.attributes = attributes
393 self.inputs = TosaSerializer.toList(inputs)
394 self.outputs = TosaSerializer.toList(outputs)
395 self.quantInfo = quantInfo
396
397 def __str__(self):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800398 str = "Op {}\n----\n".format(self.op)
Eric Kunzee5e26762020-10-13 16:11:07 -0700399
400 for i in self.inputs:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800401 str = str + " Input: {}\n".format(i)
Eric Kunzee5e26762020-10-13 16:11:07 -0700402 for o in self.outputs:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800403 str = str + " Output: {}\n".format(o)
Eric Kunzee5e26762020-10-13 16:11:07 -0700404
405 return str
406
407 def serialize(self, builder):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800408 fb_inputs = TosaSerializer.serializeStrVec(
409 builder, self.inputs, TosaOperator.TosaOperatorStartInputsVector
410 )
411 fb_outputs = TosaSerializer.serializeStrVec(
412 builder, self.outputs, TosaOperator.TosaOperatorStartOutputsVector
413 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700414 # Need to serialize quant_info and attributes enums still
415 if self.attributes is not None:
416 fb_attributes = self.attributes.serialize(builder)
417
418 if self.quantInfo is not None:
419 fb_qinfo = self.quantInfo.serialize(builder)
420
421 TosaOperator.TosaOperatorStart(builder)
422 TosaOperator.TosaOperatorAddOp(builder, self.op)
423 TosaOperator.TosaOperatorAddInputs(builder, fb_inputs)
424 TosaOperator.TosaOperatorAddOutputs(builder, fb_outputs)
425 if self.attributes is not None:
426 TosaOperator.TosaOperatorAddAttributeType(builder, self.attributes.utype)
427 TosaOperator.TosaOperatorAddAttribute(builder, fb_attributes)
428 if self.quantInfo is not None:
429 TosaOperator.TosaOperatorAddQuantInfoType(builder, self.quantInfo.utype)
430 TosaOperator.TosaOperatorAddQuantInfo(builder, fb_qinfo)
431
432 return TosaOperator.TosaOperatorEnd(builder)
433
Kevin Cheng550ccc52021-03-03 11:21:43 -0800434
Eric Kunzee5e26762020-10-13 16:11:07 -0700435class TosaSerializerBasicBlock:
436 def __init__(self, name):
437 self.name = name
438 self.operators = []
439
440 # Dict assures uniqueness, but allows us to look up by name
441 self.tensors = dict()
442
443 self.inputs = []
444 self.outputs = []
445
Kevin Cheng550ccc52021-03-03 11:21:43 -0800446 def addTensor(
447 self,
448 name,
449 shape,
450 dtype,
451 filename=None,
452 placeholderFilename=None,
453 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700454 try:
455 # Someone already added this tensor.
Eric Kunzee5e26762020-10-13 16:11:07 -0700456 tens = self.tensors[name]
Eric Kunzee5e26762020-10-13 16:11:07 -0700457 except KeyError:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800458 self.tensors[name] = TosaSerializerTensor(
459 name, shape, dtype, filename, placeholderFilename
460 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700461
462 return self.tensors[name]
463
464 def addInput(self, name):
465 self.inputs.append(name)
466
467 def addOutput(self, name):
468 self.outputs.append(name)
469
Kevin Cheng550ccc52021-03-03 11:21:43 -0800470 def addOperator(self, op, inputs, outputs, attributes=None, quant_info=None):
471 self.operators.append(
472 TosaSerializerOperator(op, inputs, outputs, attributes, quant_info)
473 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700474
475 def serialize(self, builder):
476 fb_name = builder.CreateString(self.name)
Kevin Cheng550ccc52021-03-03 11:21:43 -0800477 fbv_inputs = TosaSerializer.serializeStrVec(
478 builder, list(self.inputs), TosaBasicBlock.TosaBasicBlockStartInputsVector
479 )
480 fbv_outputs = TosaSerializer.serializeStrVec(
481 builder, list(self.outputs), TosaBasicBlock.TosaBasicBlockStartOutputsVector
482 )
483 fbv_tensors = TosaSerializer.serializeObjVec(
484 builder,
485 list(self.tensors.values()),
486 TosaBasicBlock.TosaBasicBlockStartTensorsVector,
487 )
488 fbv_operators = TosaSerializer.serializeObjVec(
489 builder, self.operators, TosaBasicBlock.TosaBasicBlockStartOperatorsVector
490 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700491
492 TosaBasicBlock.TosaBasicBlockStart(builder)
493 TosaBasicBlock.TosaBasicBlockAddName(builder, fb_name)
494 TosaBasicBlock.TosaBasicBlockAddInputs(builder, fbv_inputs)
495 TosaBasicBlock.TosaBasicBlockAddOutputs(builder, fbv_outputs)
496 TosaBasicBlock.TosaBasicBlockAddTensors(builder, fbv_tensors)
497 TosaBasicBlock.TosaBasicBlockAddOperators(builder, fbv_operators)
498 return TosaBasicBlock.TosaBasicBlockEnd(builder)
499
Kevin Cheng550ccc52021-03-03 11:21:43 -0800500
Eric Kunzee5e26762020-10-13 16:11:07 -0700501@unique
502class TensorDir(IntEnum):
503 PLACEHOLDER = 0
504 CONST = 1
505 INTERMEDIATE = 2
506 RESULT = 3
507
Kevin Cheng550ccc52021-03-03 11:21:43 -0800508
Eric Kunzee5e26762020-10-13 16:11:07 -0700509class TosaSerializer:
510 def __init__(self, pathPrefix):
511
512 # Get the global TOSA version if not already defined
513 try:
514 TOSA_VERSION
515 except NameError:
516 TosaSerializer.setTosaVersion()
517
518 self.builder = flatbuffers.Builder(0)
519
520 self.basicBlocks = []
Kevin Cheng550ccc52021-03-03 11:21:43 -0800521 self.startBasicBlock("main")
Eric Kunzee5e26762020-10-13 16:11:07 -0700522 self.pathPrefix = pathPrefix
523
524 # Indicies used for adding/naming tensors
525 self.currInputIdx = 0
526 self.currConstIdx = 0
527 self.currLayerIdx = 1
528 self.currResultIdx = 0
529
530 # Is this an illegal test that is expected to fail?
531 self.expectedFailure = False
Kevin Cheng550ccc52021-03-03 11:21:43 -0800532 self.expectedFailureDesc = ""
Eric Kunzee5e26762020-10-13 16:11:07 -0700533
534 def __str__(self):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800535 str = ""
Eric Kunzee5e26762020-10-13 16:11:07 -0700536 for bb in self.basicBlocks:
537 str = str + bb.__str__()
538 return str
539
Kevin Cheng550ccc52021-03-03 11:21:43 -0800540 def addPlaceholder(self, shape, dtype, vals):
Eric Kunzee5e26762020-10-13 16:11:07 -0700541 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800542 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700543
Kevin Cheng550ccc52021-03-03 11:21:43 -0800544 name = "input-{}".format(self.currInputIdx)
545 filename = "{}.npy".format(name)
Eric Kunzee5e26762020-10-13 16:11:07 -0700546 self.currInputIdx = self.currInputIdx + 1
547
Kevin Cheng550ccc52021-03-03 11:21:43 -0800548 tens = self.currBasicBlock.addTensor(name, shape, dtype, None, filename)
Eric Kunzee5e26762020-10-13 16:11:07 -0700549 # This is always an input to the block
550 self.currBasicBlock.addInput(name)
Eric Kunzee5e26762020-10-13 16:11:07 -0700551
552 if vals is not None:
553 np.save(os.path.join(self.pathPrefix, filename), vals, False)
554
555 return tens
556
Kevin Cheng550ccc52021-03-03 11:21:43 -0800557 def addConst(self, shape, dtype, vals):
Eric Kunzee5e26762020-10-13 16:11:07 -0700558 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800559 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700560
Kevin Cheng550ccc52021-03-03 11:21:43 -0800561 name = "const-{}".format(self.currInputIdx)
562 filename = "{}.npy".format(name)
Eric Kunzee5e26762020-10-13 16:11:07 -0700563 self.currInputIdx = self.currInputIdx + 1
564
Kevin Cheng550ccc52021-03-03 11:21:43 -0800565 tens = self.currBasicBlock.addTensor(name, shape, dtype, filename)
Eric Kunzee5e26762020-10-13 16:11:07 -0700566 # Add the operator now
567 self.currBasicBlock.addOperator(tosa.Op.Op().CONST, [], name)
568
569 if vals is not None:
570 np.save(os.path.join(self.pathPrefix, filename), vals, False)
571 return tens
572
Kevin Cheng550ccc52021-03-03 11:21:43 -0800573 def addIntermediate(self, shape, dtype):
Eric Kunzee5e26762020-10-13 16:11:07 -0700574
575 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800576 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700577
Kevin Cheng550ccc52021-03-03 11:21:43 -0800578 name = "layer-{}".format(self.currLayerIdx)
579 filename = None # No file, so no filename
Eric Kunzee5e26762020-10-13 16:11:07 -0700580 self.currLayerIdx = self.currLayerIdx + 1
581
Kevin Cheng550ccc52021-03-03 11:21:43 -0800582 tens = self.currBasicBlock.addTensor(name, shape, dtype, filename)
Eric Kunzee5e26762020-10-13 16:11:07 -0700583
584 return tens
585
586 def addInputTensor(self, tensor):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800587 self.currBasicBlock.addTensor(tensor.name, tensor.shape, tensor.dtype)
Eric Kunzee5e26762020-10-13 16:11:07 -0700588 self.currBasicBlock.addInput(tensor.name)
589
590 def addOutputTensor(self, tensor):
591 self.currBasicBlock.addOutput(tensor.name)
592
Kevin Cheng550ccc52021-03-03 11:21:43 -0800593 def addOutput(self, shape, dtype):
Eric Kunzee5e26762020-10-13 16:11:07 -0700594 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800595 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700596
Kevin Cheng550ccc52021-03-03 11:21:43 -0800597 name = "result-{}".format(self.currResultIdx)
Eric Kunzee5e26762020-10-13 16:11:07 -0700598 self.currResultIdx = self.currResultIdx + 1
599
Kevin Cheng550ccc52021-03-03 11:21:43 -0800600 tens = self.currBasicBlock.addTensor(name, shape, dtype, None)
Eric Kunzee5e26762020-10-13 16:11:07 -0700601 self.currBasicBlock.addOutput(name)
602 return tens
603
Kevin Cheng550ccc52021-03-03 11:21:43 -0800604 def addOperator(self, op, inputs, outputs, attributes=None, quant_info=None):
Eric Kunzee5e26762020-10-13 16:11:07 -0700605
Kevin Cheng14d7f7a2021-05-12 10:44:49 -0700606 if op == tosa.Op.Op().CONST:
607 raise Exception("Use addConstTensor() to add CONST ops")
Eric Kunzee5e26762020-10-13 16:11:07 -0700608
Kevin Cheng550ccc52021-03-03 11:21:43 -0800609 return self.currBasicBlock.addOperator(
610 op, inputs, outputs, attributes, quant_info
611 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700612
Kevin Cheng550ccc52021-03-03 11:21:43 -0800613 def setExpectedFailure(self, desc="", val=True):
Eric Kunzee5e26762020-10-13 16:11:07 -0700614
Eric Kunzee5e26762020-10-13 16:11:07 -0700615 self.expectedFailure = val
616 self.expectedFailureDesc = desc
617
618 def serialize(self):
619
620 builder = self.builder
621
622 Version.VersionStart(builder)
623 Version.VersionAdd_major(builder, TOSA_VERSION[0])
624 Version.VersionAdd_minor(builder, TOSA_VERSION[1])
625 Version.VersionAdd_patch(builder, TOSA_VERSION[2])
626 Version.VersionAdd_experimental(builder, TOSA_VERSION[3])
627 version = Version.VersionEnd(builder)
628
Kevin Cheng550ccc52021-03-03 11:21:43 -0800629 fbv_bb = TosaSerializer.serializeObjVec(
630 builder, self.basicBlocks, TosaGraph.TosaGraphStartBlocksVector
631 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700632
633 TosaGraph.TosaGraphStart(builder)
634 TosaGraph.TosaGraphAddVersion(builder, version)
635 TosaGraph.TosaGraphAddBlocks(builder, fbv_bb)
636 graph = TosaGraph.TosaGraphEnd(builder)
637
638 self.builder.Finish(graph)
639 return self.builder.Output()
640
641 def writeJson(self, tosa_filename):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800642 """Write a json test file so that it is fairly easy to pick up the test
643 and generate commands for third party tool"""
Eric Kunzee5e26762020-10-13 16:11:07 -0700644 test_desc = dict()
645
Kevin Cheng550ccc52021-03-03 11:21:43 -0800646 test_desc["tosa_file"] = tosa_filename
Eric Kunzee5e26762020-10-13 16:11:07 -0700647 ifm_name = []
648 ifm_shape = []
649 ifm_file = []
650 ofm_name = []
651 ofm_file = []
652 ofm_shape = []
653
654 for b in self.basicBlocks:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800655 if b.name == "main":
Eric Kunzee5e26762020-10-13 16:11:07 -0700656 for i in b.inputs:
657 ifm_name.append(i)
658 ifm_shape.append(b.tensors[i].shape)
659 ifm_file.append(b.tensors[i].placeholderFilename)
660 for o in b.outputs:
661 ofm_name.append(o)
662 ofm_shape.append(b.tensors[o].shape)
663 # Make up an OFM filename here. One isn't generated until the reference tool is
664 # run, so any name is a good name
Kevin Cheng550ccc52021-03-03 11:21:43 -0800665 ofm_file.append("ref-{}.npy".format(o))
Eric Kunzee5e26762020-10-13 16:11:07 -0700666
Kevin Cheng550ccc52021-03-03 11:21:43 -0800667 test_desc["ifm_placeholder"] = ifm_name
668 test_desc["ifm_file"] = ifm_file
669 test_desc["ifm_shape"] = ifm_shape
670 test_desc["ofm_name"] = ofm_name
671 test_desc["ofm_shape"] = ofm_shape
672 test_desc["ofm_file"] = ofm_file
673 test_desc["expected_failure"] = self.expectedFailure
Eric Kunzee5e26762020-10-13 16:11:07 -0700674 if self.expectedFailureDesc:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800675 test_desc["expected_failure_desc"] = self.expectedFailureDesc
Eric Kunzee5e26762020-10-13 16:11:07 -0700676
Kevin Cheng550ccc52021-03-03 11:21:43 -0800677 return json.dumps(test_desc, indent=" ")
Eric Kunzee5e26762020-10-13 16:11:07 -0700678
679 def startBasicBlock(self, name):
680 self.currBasicBlock = TosaSerializerBasicBlock(name)
681 self.basicBlocks.append(self.currBasicBlock)
682
683 @staticmethod
684 def serializeStrVec(builder, vec, start_fcn):
685 fb_strs = [builder.CreateString(i) for i in vec]
686 start_fcn(builder, len(fb_strs))
687 for s in fb_strs[::-1]:
688 builder.PrependUOffsetTRelative(s)
689 return builder.EndVector(len(fb_strs))
690
691 @staticmethod
692 def serializeInt32Vec(builder, vec):
693 builder.StartVector(4, len(vec), 4)
694 for v in vec[::-1]:
695 builder.PrependInt32(v)
696 return builder.EndVector(len(vec))
697
698 @staticmethod
Kevin Cheng77d0f762020-11-24 10:26:32 -0800699 def serializeFpVec(builder, vec):
700 builder.StartVector(4, len(vec), 4)
701 for v in vec[::-1]:
702 builder.PrependFloat32(v)
703 return builder.EndVector(len(vec))
704
705 @staticmethod
Eric Kunzee5e26762020-10-13 16:11:07 -0700706 def serializeObjVec(builder, vec, start_fcn):
707 serialized_vec = []
708 for v in vec[::-1]:
709 serialized_vec.append(v.serialize(builder))
710
711 start_fcn(builder, len(vec))
712 for v in serialized_vec:
713 builder.PrependUOffsetTRelative(v)
714 return builder.EndVector(len(vec))
715
716 @staticmethod
717 def toList(val):
718 if isinstance(val, list):
719 return val
720 else:
721 return [val]
722
723 @staticmethod
724 def setTosaVersion():
725 # Create a dummy flatbuffers file with the default version information
726 # There does not appear to be a better way to get a constant from a
727 # flatbuffer schema file
728 builder = flatbuffers.Builder(0)
729 Version.VersionStart(builder)
730 ver = Version.VersionEnd(builder)
731 TosaGraph.TosaGraphStart(builder)
732 TosaGraph.TosaGraphAddVersion(builder, ver)
733 gr = TosaGraph.TosaGraphEnd(builder)
734 builder.Finish(gr)
735
736 out = builder.Output()
737
738 gr = TosaGraph.TosaGraph()
739 root = gr.GetRootAsTosaGraph(out, 0)
740
741 # Store the version as a global variable so that it only needs to be
742 # generated once per process.
743 global TOSA_VERSION
Kevin Cheng550ccc52021-03-03 11:21:43 -0800744 TOSA_VERSION = [
745 root.Version()._major(),
746 root.Version()._minor(),
747 root.Version()._patch(),
748 root.Version()._experimental(),
749 ]