blob: c4de2a22a9fa12f6d2b53bac3fc7e706d5f349dc [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
Kevin Cheng82507d72021-06-17 16:01:59 -070022import struct
Eric Kunzee5e26762020-10-13 16:11:07 -070023from enum import Enum, IntEnum, unique
Kevin Cheng550ccc52021-03-03 11:21:43 -080024from tosa import (
25 TosaGraph,
26 TosaBasicBlock,
27 TosaTensor,
28 TosaOperator,
29 DType,
30 Op,
31 ResizeMode,
32 Version,
33)
34
35# Include the ../thirdparty/serialization_lib/python directory in PYTHONPATH
36parent_dir = os.path.dirname(os.path.realpath(__file__))
37sys.path.append(
38 os.path.join(parent_dir, "..", "thirdparty", "serialization_lib", "python")
39)
Eric Kunzee5e26762020-10-13 16:11:07 -070040import tosa
Eric Kunzee5e26762020-10-13 16:11:07 -070041
42# With the way flatc generates its python types, there is no programatic way
43# to get string names for the integer types. Manually maintain a string table
44# here.
Kevin Cheng82507d72021-06-17 16:01:59 -070045DType = tosa.DType.DType()
Kevin Cheng550ccc52021-03-03 11:21:43 -080046DTypeNames = [
47 "UNKNOWN",
48 "BOOL",
49 "UINT8",
50 "INT4",
51 "INT8",
52 "INT16",
53 "INT32",
54 "INT48",
55 "FLOAT",
56]
57
Kevin Cheng82507d72021-06-17 16:01:59 -070058ByteMask = np.uint64(0xFF)
Eric Kunzee5e26762020-10-13 16:11:07 -070059
60def dtype_str_to_val(name):
61
62 for i in range(len(DTypeNames)):
63 if name.casefold() == DTypeNames[i].casefold():
64 return i
Kevin Cheng550ccc52021-03-03 11:21:43 -080065 raise Exception("Unable to parse DType name {}".format(name))
Eric Kunzee5e26762020-10-13 16:11:07 -070066
67
68class TosaSerializerUnion:
Kevin Cheng550ccc52021-03-03 11:21:43 -080069 """This class handles encapsulating and serializing union types into flatbuffers"""
70
Eric Kunzee5e26762020-10-13 16:11:07 -070071 def __init__(self):
72
73 # A tuple of the start and end functions. Set by the options constructors below
74 self.optFcns = None
75
76 # The type from the tosa.Options enumeration. Set by the options constructors below.
77 self.utype = None
78
79 # Each of these lists is a tuple of the add function and the
80 # value being added. Set by the options constructors below.
81 self.ints = []
82 self.bools = []
83 self.floats = []
84 self.strings = []
85 self.intvecs = []
Kevin Cheng77d0f762020-11-24 10:26:32 -080086 self.fpvecs = []
Eric Kunzee5e26762020-10-13 16:11:07 -070087
88 def serialize(self, builder):
89
90 # We have to build strings and vectors first
91 strList = []
92 intVecList = []
Kevin Cheng77d0f762020-11-24 10:26:32 -080093 fpVecList = []
Eric Kunzee5e26762020-10-13 16:11:07 -070094
95 for fcn, val in self.strings:
96 strList.append((fcn, builder.CreateString(val)))
97
98 for fcn, val in self.intvecs:
99 intVecList.append((fcn, TosaSerializer.serializeInt32Vec(builder, val)))
100
Kevin Cheng77d0f762020-11-24 10:26:32 -0800101 for fcn, val in self.fpvecs:
102 fpVecList.append((fcn, TosaSerializer.serializeFpVec(builder, val)))
103
Eric Kunzee5e26762020-10-13 16:11:07 -0700104 startFcn, endFcn = self.optFcns
105
106 # Then serialize the options object from the list of primitives and
107 # other serialized values
108 startFcn(builder)
109 for fcn, val in self.ints:
110 fcn(builder, val)
111
112 for fcn, val in self.bools:
113 fcn(builder, val)
114
115 for fcn, val in self.floats:
116 fcn(builder, val)
117
118 for fcn, val in strList:
119 fcn(builder, val)
120
121 for fcn, val in intVecList:
122 fcn(builder, val)
123
Kevin Cheng77d0f762020-11-24 10:26:32 -0800124 for fcn, val in fpVecList:
125 fcn(builder, val)
126
Eric Kunzee5e26762020-10-13 16:11:07 -0700127 return endFcn(builder)
128
Kevin Cheng550ccc52021-03-03 11:21:43 -0800129
Eric Kunzee5e26762020-10-13 16:11:07 -0700130class TosaSerializerAttribute(TosaSerializerUnion):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800131 """This class handles encapsulating all of the enumerated types for attributes"""
Eric Kunzee5e26762020-10-13 16:11:07 -0700132
133 def __init__(self):
134 super().__init__()
135
136 def Pool2dAttribute(self, kernel, stride, padding):
137 from tosa import Pool2dAttribute as a, Attribute
138
139 self.utype = Attribute.Attribute().Pool2dAttribute
140
141 self.optFcns = (a.Pool2dAttributeStart, a.Pool2dAttributeEnd)
Kevin Cheng550ccc52021-03-03 11:21:43 -0800142 self.intvecs.append((a.Pool2dAttributeAddPadding, padding))
143 self.intvecs.append((a.Pool2dAttributeAddKernel, kernel))
144 self.intvecs.append((a.Pool2dAttributeAddStride, stride))
Eric Kunzee5e26762020-10-13 16:11:07 -0700145
146 def Conv2dAttribute(self, padding, stride, dilation):
147 from tosa import Conv2dAttribute as a, Attribute
148
149 self.utype = Attribute.Attribute().Conv2dAttribute
150 self.optFcns = (a.Conv2dAttributeStart, a.Conv2dAttributeEnd)
151
Kevin Cheng550ccc52021-03-03 11:21:43 -0800152 self.intvecs.append((a.Conv2dAttributeAddPadding, padding))
153 self.intvecs.append((a.Conv2dAttributeAddStride, stride))
154 self.intvecs.append((a.Conv2dAttributeAddDilation, dilation))
Eric Kunzee5e26762020-10-13 16:11:07 -0700155
156 def TransposeConv2DAttribute(self, outpad, stride, dilation, output_shape):
157 from tosa import TransposeConv2dAttribute as a, Attribute
158
159 self.utype = Attribute.Attribute().TransposeConv2dAttribute
160 self.optFcns = (a.TransposeConv2dAttributeStart, a.TransposeConv2dAttributeEnd)
161
Kevin Cheng550ccc52021-03-03 11:21:43 -0800162 self.intvecs.append((a.TransposeConv2dAttributeAddOutpad, outpad))
163 self.intvecs.append((a.TransposeConv2dAttributeAddStride, stride))
164 self.intvecs.append((a.TransposeConv2dAttributeAddDilation, dilation))
165 self.intvecs.append((a.TransposeConv2dAttributeAddOutputShape, output_shape))
Eric Kunzee5e26762020-10-13 16:11:07 -0700166
167 def ReluNAttribute(self, maxint, maxfp):
168 from tosa import ReluNAttribute as a, Attribute
169
170 self.utype = Attribute.Attribute().ReluNAttribute
171 self.optFcns = (a.ReluNAttributeStart, a.ReluNAttributeEnd)
172
173 self.ints.append((a.ReluNAttributeAddMaxInt, maxint))
174 self.ints.append((a.ReluNAttributeAddMaxFp, maxfp))
175
Eric Kunzee5e26762020-10-13 16:11:07 -0700176 def AxisAttribute(self, axis):
177 from tosa import AxisAttribute as a, Attribute
178
179 self.utype = Attribute.Attribute().AxisAttribute
180 self.optFcns = (a.AxisAttributeStart, a.AxisAttributeEnd)
181
Kevin Cheng550ccc52021-03-03 11:21:43 -0800182 self.ints.append((a.AxisAttributeAddAxis, axis))
Eric Kunzee5e26762020-10-13 16:11:07 -0700183
184 def ReshapeAttribute(self, shape):
185 from tosa import ReshapeAttribute as a, Attribute
186
187 self.utype = Attribute.Attribute().ReshapeAttribute
188 self.optFcns = (a.ReshapeAttributeStart, a.ReshapeAttributeEnd)
189
Kevin Cheng550ccc52021-03-03 11:21:43 -0800190 self.intvecs.append((a.ReshapeAttributeAddShape, shape))
Eric Kunzee5e26762020-10-13 16:11:07 -0700191
192 def SliceAttribute(self, begin, size):
193 from tosa import SliceAttribute as a, Attribute
194
195 self.utype = Attribute.Attribute().SliceAttribute
196 self.optFcns = (a.SliceAttributeStart, a.SliceAttributeEnd)
197
Kevin Cheng550ccc52021-03-03 11:21:43 -0800198 self.intvecs.append((a.SliceAttributeAddBegin, begin))
199 self.intvecs.append((a.SliceAttributeAddSize, size))
Eric Kunzee5e26762020-10-13 16:11:07 -0700200
201 def TileAttribute(self, multiples):
202 from tosa import TileAttribute as a, Attribute
203
204 self.utype = Attribute.Attribute().TileAttribute
205 self.optFcns = (a.TileAttributeStart, a.TileAttributeEnd)
206
Kevin Cheng550ccc52021-03-03 11:21:43 -0800207 self.intvecs.append((a.TileAttributeAddMultiples, multiples))
Eric Kunzee5e26762020-10-13 16:11:07 -0700208
Kevin Cheng550ccc52021-03-03 11:21:43 -0800209 def ResizeAttribute(
210 self, output_size, stride, offset, shift, stride_fp, offset_fp, mode
211 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700212 from tosa import ResizeAttribute as a, Attribute
213
214 self.utype = Attribute.Attribute().ResizeAttribute
215 self.optFcns = (a.ResizeAttributeStart, a.ResizeAttributeEnd)
216
Kevin Cheng550ccc52021-03-03 11:21:43 -0800217 self.intvecs.append((a.ResizeAttributeAddOutputSize, output_size))
218 self.intvecs.append((a.ResizeAttributeAddStride, stride))
219 self.intvecs.append((a.ResizeAttributeAddOffset, offset))
220 self.ints.append((a.ResizeAttributeAddShift, shift))
221 self.fpvecs.append((a.ResizeAttributeAddStrideFp, stride_fp))
222 self.fpvecs.append((a.ResizeAttributeAddOffsetFp, offset_fp))
223 self.ints.append((a.ResizeAttributeAddMode, mode))
Eric Kunzee5e26762020-10-13 16:11:07 -0700224
225 def ClampAttribute(self, minint, maxint, minfp, maxfp):
226 from tosa import ClampAttribute as a, Attribute
227
228 self.utype = Attribute.Attribute().ClampAttribute
229 self.optFcns = (a.ClampAttributeStart, a.ClampAttributeEnd)
230
Kevin Cheng550ccc52021-03-03 11:21:43 -0800231 self.ints.append((a.ClampAttributeAddMinInt, minint))
232 self.ints.append((a.ClampAttributeAddMaxInt, maxint))
Eric Kunzee5e26762020-10-13 16:11:07 -0700233
Kevin Cheng550ccc52021-03-03 11:21:43 -0800234 self.ints.append((a.ClampAttributeAddMinFp, minfp))
235 self.ints.append((a.ClampAttributeAddMaxFp, maxfp))
Eric Kunzee5e26762020-10-13 16:11:07 -0700236
Kevin Cheng550ccc52021-03-03 11:21:43 -0800237 def RescaleAttribute(
238 self, input_zp, output_zp, multiplier, shift, scale32, double_round, per_channel
239 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700240 from tosa import RescaleAttribute as a, Attribute
241
242 self.utype = Attribute.Attribute().RescaleAttribute
243 self.optFcns = (a.RescaleAttributeStart, a.RescaleAttributeEnd)
244
Kevin Cheng550ccc52021-03-03 11:21:43 -0800245 self.ints.append((a.RescaleAttributeAddInputZp, input_zp))
246 self.ints.append((a.RescaleAttributeAddOutputZp, output_zp))
247 self.intvecs.append((a.RescaleAttributeAddMultiplier, multiplier))
248 self.intvecs.append((a.RescaleAttributeAddShift, shift))
249 self.bools.append((a.RescaleAttributeAddScale32, scale32))
250 self.bools.append((a.RescaleAttributeAddDoubleRound, double_round))
251 self.bools.append((a.RescaleAttributeAddPerChannel, per_channel))
Eric Kunzee5e26762020-10-13 16:11:07 -0700252
Kevin Chengaee1fac2020-11-11 13:54:06 -0800253 def MulAttribute(self, shift):
254 from tosa import MulAttribute as a, Attribute
255
256 self.utype = Attribute.Attribute().MulAttribute
257 self.optFcns = (a.MulAttributeStart, a.MulAttributeEnd)
258
Kevin Cheng550ccc52021-03-03 11:21:43 -0800259 self.ints.append((a.MulAttributeAddShift, shift))
Kevin Chengaee1fac2020-11-11 13:54:06 -0800260
261 def ArithmeticRightShiftAttribute(self, round):
262 from tosa import ArithmeticRightShiftAttribute as a, Attribute
263
264 self.utype = Attribute.Attribute().ArithmeticRightShiftAttribute
Kevin Cheng550ccc52021-03-03 11:21:43 -0800265 self.optFcns = (
266 a.ArithmeticRightShiftAttributeStart,
267 a.ArithmeticRightShiftAttributeEnd,
268 )
Kevin Chengaee1fac2020-11-11 13:54:06 -0800269
Kevin Cheng550ccc52021-03-03 11:21:43 -0800270 self.bools.append((a.ArithmeticRightShiftAttributeAddRound, round))
Kevin Chengaee1fac2020-11-11 13:54:06 -0800271
Eric Kunzee5e26762020-10-13 16:11:07 -0700272 def CustomAttribute(self, identifier):
273 from tosa import CustomAttribute as a, Attribute
274
275 self.utype = Attribute.Attribute().CustomAttribute
276 self.optFcns = (a.CustomAttributeStart, a.CustomAttributeEnd)
277
Kevin Cheng550ccc52021-03-03 11:21:43 -0800278 self.strings.append((a.CustomAttributeAddIdentifier, identifier))
Eric Kunzee5e26762020-10-13 16:11:07 -0700279
280 def CondIfAttribute(self, then_branch, else_branch):
281 from tosa import CondIfAttribute as a, Attribute
282
283 self.utype = Attribute.Attribute().CondIfAttribute
284 self.optFcns = (a.CondIfAttributeStart, a.CondIfAttributeEnd)
285
Kevin Cheng550ccc52021-03-03 11:21:43 -0800286 self.strings.append((a.CondIfAttributeAddThenBranch, then_branch))
287 self.strings.append((a.CondIfAttributeAddElseBranch, else_branch))
Eric Kunzee5e26762020-10-13 16:11:07 -0700288
289 def WhileLoopAttribute(self, cond_branch, body_branch):
290 from tosa import WhileLoopAttribute as a, Attribute
291
292 self.utype = Attribute.Attribute().WhileLoopAttribute
293 self.optFcns = (a.WhileLoopAttributeStart, a.WhileLoopAttributeEnd)
294
Kevin Cheng550ccc52021-03-03 11:21:43 -0800295 self.strings.append((a.WhileLoopAttributeAddCondBranch, cond_branch))
296 self.strings.append((a.WhileLoopAttributeAddBodyBranch, body_branch))
297
Eric Kunzee5e26762020-10-13 16:11:07 -0700298
299class TosaSerializerQuantInfo(TosaSerializerUnion):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800300 """This class handles encapsulating all of the enumerated types for quantinfo types"""
301
Eric Kunzee5e26762020-10-13 16:11:07 -0700302 def __init__(self):
303 super().__init__()
304
305 def ConvQuantInfo(self, input_zp, weight_zp):
306 from tosa import ConvQuantInfo as q, QuantInfo
307
308 self.utype = QuantInfo.QuantInfo().ConvQuantInfo
309 self.optFcns = (q.ConvQuantInfoStart, q.ConvQuantInfoEnd)
310 self.ints.append((q.ConvQuantInfoAddInputZp, input_zp))
311 self.ints.append((q.ConvQuantInfoAddWeightZp, weight_zp))
312
313 def UnaryQuantInfo(self, input_zp, output_zp):
314 from tosa import UnaryQuantInfo as q, QuantInfo
315
316 self.utype = QuantInfo.QuantInfo().UnaryQuantInfo
317 self.optFcns = (q.UnaryQuantInfoStart, q.UnaryQuantInfoEnd)
318 self.ints.append((q.UnaryQuantInfoAddInputZp, input_zp))
319 self.ints.append((q.UnaryQuantInfoAddOutputZp, output_zp))
320
321 def MatMulQuantInfo(self, a_zp, b_zp):
322 from tosa import MatMulQuantInfo as q, QuantInfo
323
324 self.utype = QuantInfo.QuantInfo().MatMulQuantInfo
325 self.optFcns = (q.MatMulQuantInfoStart, q.MatMulQuantInfoEnd)
326 self.ints.append((q.MatMulQuantInfoAddAZp, a_zp))
327 self.ints.append((q.MatMulQuantInfoAddBZp, b_zp))
328
329 def PadQuantInfo(self, input_zp):
330 from tosa import PadQuantInfo as q, QuantInfo
331
332 self.utype = QuantInfo.QuantInfo().PadQuantInfo
333 self.optFcns = (q.PadQuantInfoStart, q.PadQuantInfoEnd)
334 self.ints.append((q.PadQuantInfoAddInputZp, input_zp))
335
Kevin Cheng550ccc52021-03-03 11:21:43 -0800336
Eric Kunzee5e26762020-10-13 16:11:07 -0700337class TosaSerializerTensor:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800338 def __init__(
339 self,
340 name,
341 shape,
342 dtype,
Kevin Cheng82507d72021-06-17 16:01:59 -0700343 data=None,
Kevin Cheng550ccc52021-03-03 11:21:43 -0800344 placeholderFilename=None,
345 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700346 self.name = name
347
348 if isinstance(shape, np.ndarray):
349 shape = shape.astype(int).tolist()
350 shape = list(map(int, shape))
351
352 self.shape = shape
353 self.dtype = dtype
Eric Kunzee5e26762020-10-13 16:11:07 -0700354
Kevin Cheng82507d72021-06-17 16:01:59 -0700355 if isinstance(data, np.ndarray):
356 data = data.flatten().astype(int).tolist()
357 data = list(map(int, data))
358 self.data = data
359 else:
360 self.data = None
Eric Kunzee5e26762020-10-13 16:11:07 -0700361
362 # Filename for placeholder tensors. These get generated by the test generation
363 # process and are written to disk, but are considered input tensors by the network
364 # so they do not appear in the TOSA serialiazation. However, if we want to form a unit
365 # test around these input tensors, we can get the filename from here.
366 self.placeholderFilename = placeholderFilename
367
368 def __str__(self):
Kevin Cheng82507d72021-06-17 16:01:59 -0700369 str = "TosaSerializerTensor name: {} shape: {} dtype: {}".format(
Kevin Cheng550ccc52021-03-03 11:21:43 -0800370 self.name,
371 self.shape,
372 DTypeNames[self.dtype],
Kevin Cheng550ccc52021-03-03 11:21:43 -0800373 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700374 return str
375
Eric Kunzee5e26762020-10-13 16:11:07 -0700376 def setDtype(self, dtype):
377 self.dtype = dtype
378
Eric Kunzee5e26762020-10-13 16:11:07 -0700379 def serialize(self, builder):
380 fb_name = builder.CreateString(self.name)
Eric Kunzee5e26762020-10-13 16:11:07 -0700381 fb_shapes = TosaSerializer.serializeInt32Vec(builder, self.shape)
Kevin Cheng82507d72021-06-17 16:01:59 -0700382 if self.data:
383 u8_data = list()
384 # little endianess
385 if self.dtype == DType.BOOL:
386 for val in self.data:
387 val_u8 = np.uint8(val)
388 u8_data.append(val_u8)
389 elif self.dtype == DType.INT8:
390 for val in self.data:
391 val_u8 = np.uint8(val)
392 u8_data.append(val_u8)
393 elif self.dtype == DType.INT16:
394 for val in self.data:
395 val_u16 = np.uint16(val)
396 b0 = val_u16 & ByteMask
397 b1 = (val_u16 >> np.uint16(8)) & ByteMask
398 u8_data.extend([b0, b1])
399 elif self.dtype == DType.INT32:
400 for val in self.data:
401 val_u32 = np.uint32(val)
402 b0 = val_u32 & ByteMask
403 b1 = (val_u32 >> np.uint32(8)) & ByteMask
404 b2 = (val_u32 >> np.uint32(16)) & ByteMask
405 b3 = (val_u32 >> np.uint32(32)) & ByteMask
406 u8_data.extend([b0, b1, b2, b3])
407 elif self.dtype == DType.INT48:
408 for val in self.data:
409 val_u64 = np.uint64(val)
410 b0 = val_u64 & ByteMask
411 b1 = (val_u64 >> np.uint64(8)) & ByteMask
412 b2 = (val_u64 >> np.uint64(16)) & ByteMask
413 b3 = (val_u64 >> np.uint64(24)) & ByteMask
414 b4 = (val_u64 >> np.uint64(32)) & ByteMask
415 b5 = (val_u64 >> np.uint64(40)) & ByteMask
416 u8_data.extend([b0, b1, b2, b3, b4, b5])
417 elif self.dtype == DType.FLOAT:
418 for val in self.data:
419 b = struct.pack('!f', val)
420 u8_data.extend([b[3], b[2], b[1], b[0]])
421 else:
422 raise Exception("unsupported data type {}".format(DTypeNames[self.dtype]))
423 fb_data = TosaSerializer.serializeUint8Vec(builder, u8_data)
Eric Kunzee5e26762020-10-13 16:11:07 -0700424
425 TosaTensor.TosaTensorStart(builder)
426 TosaTensor.TosaTensorAddName(builder, fb_name)
427 TosaTensor.TosaTensorAddShape(builder, fb_shapes)
428 TosaTensor.TosaTensorAddType(builder, self.dtype)
Kevin Cheng82507d72021-06-17 16:01:59 -0700429 if self.data:
430 TosaTensor.TosaTensorAddData(builder, fb_data)
Eric Kunzee5e26762020-10-13 16:11:07 -0700431
432 return TosaTensor.TosaTensorEnd(builder)
433
Kevin Cheng550ccc52021-03-03 11:21:43 -0800434
Eric Kunzee5e26762020-10-13 16:11:07 -0700435class TosaSerializerOperator:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800436 def __init__(self, op, inputs, outputs, attributes=None, quantInfo=None):
Eric Kunzee5e26762020-10-13 16:11:07 -0700437 self.op = op
438 self.attributes = attributes
439 self.inputs = TosaSerializer.toList(inputs)
440 self.outputs = TosaSerializer.toList(outputs)
441 self.quantInfo = quantInfo
442
443 def __str__(self):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800444 str = "Op {}\n----\n".format(self.op)
Eric Kunzee5e26762020-10-13 16:11:07 -0700445
446 for i in self.inputs:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800447 str = str + " Input: {}\n".format(i)
Eric Kunzee5e26762020-10-13 16:11:07 -0700448 for o in self.outputs:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800449 str = str + " Output: {}\n".format(o)
Eric Kunzee5e26762020-10-13 16:11:07 -0700450
451 return str
452
453 def serialize(self, builder):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800454 fb_inputs = TosaSerializer.serializeStrVec(
455 builder, self.inputs, TosaOperator.TosaOperatorStartInputsVector
456 )
457 fb_outputs = TosaSerializer.serializeStrVec(
458 builder, self.outputs, TosaOperator.TosaOperatorStartOutputsVector
459 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700460 # Need to serialize quant_info and attributes enums still
461 if self.attributes is not None:
462 fb_attributes = self.attributes.serialize(builder)
463
464 if self.quantInfo is not None:
465 fb_qinfo = self.quantInfo.serialize(builder)
466
467 TosaOperator.TosaOperatorStart(builder)
468 TosaOperator.TosaOperatorAddOp(builder, self.op)
469 TosaOperator.TosaOperatorAddInputs(builder, fb_inputs)
470 TosaOperator.TosaOperatorAddOutputs(builder, fb_outputs)
471 if self.attributes is not None:
472 TosaOperator.TosaOperatorAddAttributeType(builder, self.attributes.utype)
473 TosaOperator.TosaOperatorAddAttribute(builder, fb_attributes)
474 if self.quantInfo is not None:
475 TosaOperator.TosaOperatorAddQuantInfoType(builder, self.quantInfo.utype)
476 TosaOperator.TosaOperatorAddQuantInfo(builder, fb_qinfo)
477
478 return TosaOperator.TosaOperatorEnd(builder)
479
Kevin Cheng550ccc52021-03-03 11:21:43 -0800480
Eric Kunzee5e26762020-10-13 16:11:07 -0700481class TosaSerializerBasicBlock:
482 def __init__(self, name):
483 self.name = name
484 self.operators = []
485
486 # Dict assures uniqueness, but allows us to look up by name
487 self.tensors = dict()
488
489 self.inputs = []
490 self.outputs = []
491
Kevin Cheng550ccc52021-03-03 11:21:43 -0800492 def addTensor(
493 self,
494 name,
495 shape,
496 dtype,
Kevin Cheng82507d72021-06-17 16:01:59 -0700497 data=None,
Kevin Cheng550ccc52021-03-03 11:21:43 -0800498 placeholderFilename=None,
499 ):
Eric Kunzee5e26762020-10-13 16:11:07 -0700500 try:
501 # Someone already added this tensor.
Eric Kunzee5e26762020-10-13 16:11:07 -0700502 tens = self.tensors[name]
Eric Kunzee5e26762020-10-13 16:11:07 -0700503 except KeyError:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800504 self.tensors[name] = TosaSerializerTensor(
Kevin Cheng82507d72021-06-17 16:01:59 -0700505 name, shape, dtype, data, placeholderFilename
Kevin Cheng550ccc52021-03-03 11:21:43 -0800506 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700507
508 return self.tensors[name]
509
510 def addInput(self, name):
511 self.inputs.append(name)
512
513 def addOutput(self, name):
514 self.outputs.append(name)
515
Kevin Cheng550ccc52021-03-03 11:21:43 -0800516 def addOperator(self, op, inputs, outputs, attributes=None, quant_info=None):
517 self.operators.append(
518 TosaSerializerOperator(op, inputs, outputs, attributes, quant_info)
519 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700520
521 def serialize(self, builder):
522 fb_name = builder.CreateString(self.name)
Kevin Cheng550ccc52021-03-03 11:21:43 -0800523 fbv_inputs = TosaSerializer.serializeStrVec(
524 builder, list(self.inputs), TosaBasicBlock.TosaBasicBlockStartInputsVector
525 )
526 fbv_outputs = TosaSerializer.serializeStrVec(
527 builder, list(self.outputs), TosaBasicBlock.TosaBasicBlockStartOutputsVector
528 )
529 fbv_tensors = TosaSerializer.serializeObjVec(
530 builder,
531 list(self.tensors.values()),
532 TosaBasicBlock.TosaBasicBlockStartTensorsVector,
533 )
534 fbv_operators = TosaSerializer.serializeObjVec(
535 builder, self.operators, TosaBasicBlock.TosaBasicBlockStartOperatorsVector
536 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700537
538 TosaBasicBlock.TosaBasicBlockStart(builder)
539 TosaBasicBlock.TosaBasicBlockAddName(builder, fb_name)
540 TosaBasicBlock.TosaBasicBlockAddInputs(builder, fbv_inputs)
541 TosaBasicBlock.TosaBasicBlockAddOutputs(builder, fbv_outputs)
542 TosaBasicBlock.TosaBasicBlockAddTensors(builder, fbv_tensors)
543 TosaBasicBlock.TosaBasicBlockAddOperators(builder, fbv_operators)
544 return TosaBasicBlock.TosaBasicBlockEnd(builder)
545
Kevin Cheng550ccc52021-03-03 11:21:43 -0800546
Eric Kunzee5e26762020-10-13 16:11:07 -0700547@unique
548class TensorDir(IntEnum):
549 PLACEHOLDER = 0
550 CONST = 1
551 INTERMEDIATE = 2
552 RESULT = 3
553
Kevin Cheng550ccc52021-03-03 11:21:43 -0800554
Eric Kunzee5e26762020-10-13 16:11:07 -0700555class TosaSerializer:
556 def __init__(self, pathPrefix):
557
558 # Get the global TOSA version if not already defined
559 try:
560 TOSA_VERSION
561 except NameError:
562 TosaSerializer.setTosaVersion()
563
564 self.builder = flatbuffers.Builder(0)
565
566 self.basicBlocks = []
Kevin Cheng550ccc52021-03-03 11:21:43 -0800567 self.startBasicBlock("main")
Eric Kunzee5e26762020-10-13 16:11:07 -0700568 self.pathPrefix = pathPrefix
569
570 # Indicies used for adding/naming tensors
571 self.currInputIdx = 0
572 self.currConstIdx = 0
573 self.currLayerIdx = 1
574 self.currResultIdx = 0
575
576 # Is this an illegal test that is expected to fail?
577 self.expectedFailure = False
Kevin Cheng550ccc52021-03-03 11:21:43 -0800578 self.expectedFailureDesc = ""
Eric Kunzee5e26762020-10-13 16:11:07 -0700579
580 def __str__(self):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800581 str = ""
Eric Kunzee5e26762020-10-13 16:11:07 -0700582 for bb in self.basicBlocks:
583 str = str + bb.__str__()
584 return str
585
Kevin Cheng550ccc52021-03-03 11:21:43 -0800586 def addPlaceholder(self, shape, dtype, vals):
Eric Kunzee5e26762020-10-13 16:11:07 -0700587 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800588 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700589
Kevin Cheng550ccc52021-03-03 11:21:43 -0800590 name = "input-{}".format(self.currInputIdx)
591 filename = "{}.npy".format(name)
Eric Kunzee5e26762020-10-13 16:11:07 -0700592 self.currInputIdx = self.currInputIdx + 1
593
Kevin Cheng550ccc52021-03-03 11:21:43 -0800594 tens = self.currBasicBlock.addTensor(name, shape, dtype, None, filename)
Eric Kunzee5e26762020-10-13 16:11:07 -0700595 # This is always an input to the block
596 self.currBasicBlock.addInput(name)
Eric Kunzee5e26762020-10-13 16:11:07 -0700597
598 if vals is not None:
599 np.save(os.path.join(self.pathPrefix, filename), vals, False)
600
601 return tens
602
Kevin Cheng550ccc52021-03-03 11:21:43 -0800603 def addConst(self, shape, dtype, vals):
Eric Kunzee5e26762020-10-13 16:11:07 -0700604 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800605 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700606
Kevin Cheng550ccc52021-03-03 11:21:43 -0800607 name = "const-{}".format(self.currInputIdx)
608 filename = "{}.npy".format(name)
Eric Kunzee5e26762020-10-13 16:11:07 -0700609 self.currInputIdx = self.currInputIdx + 1
610
Kevin Cheng82507d72021-06-17 16:01:59 -0700611 tens = self.currBasicBlock.addTensor(name, shape, dtype, vals)
Eric Kunzee5e26762020-10-13 16:11:07 -0700612 # Add the operator now
613 self.currBasicBlock.addOperator(tosa.Op.Op().CONST, [], name)
614
Eric Kunzee5e26762020-10-13 16:11:07 -0700615 return tens
616
Kevin Cheng550ccc52021-03-03 11:21:43 -0800617 def addIntermediate(self, shape, dtype):
Eric Kunzee5e26762020-10-13 16:11:07 -0700618
619 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800620 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700621
Kevin Cheng550ccc52021-03-03 11:21:43 -0800622 name = "layer-{}".format(self.currLayerIdx)
Eric Kunzee5e26762020-10-13 16:11:07 -0700623 self.currLayerIdx = self.currLayerIdx + 1
624
Kevin Cheng82507d72021-06-17 16:01:59 -0700625 tens = self.currBasicBlock.addTensor(name, shape, dtype, None)
Eric Kunzee5e26762020-10-13 16:11:07 -0700626
627 return tens
628
629 def addInputTensor(self, tensor):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800630 self.currBasicBlock.addTensor(tensor.name, tensor.shape, tensor.dtype)
Eric Kunzee5e26762020-10-13 16:11:07 -0700631 self.currBasicBlock.addInput(tensor.name)
632
633 def addOutputTensor(self, tensor):
634 self.currBasicBlock.addOutput(tensor.name)
635
Kevin Cheng550ccc52021-03-03 11:21:43 -0800636 def addOutput(self, shape, dtype):
Eric Kunzee5e26762020-10-13 16:11:07 -0700637 if not self.currBasicBlock:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800638 raise Exception("addTensor called without valid basic block")
Eric Kunzee5e26762020-10-13 16:11:07 -0700639
Kevin Cheng550ccc52021-03-03 11:21:43 -0800640 name = "result-{}".format(self.currResultIdx)
Eric Kunzee5e26762020-10-13 16:11:07 -0700641 self.currResultIdx = self.currResultIdx + 1
642
Kevin Cheng550ccc52021-03-03 11:21:43 -0800643 tens = self.currBasicBlock.addTensor(name, shape, dtype, None)
Eric Kunzee5e26762020-10-13 16:11:07 -0700644 self.currBasicBlock.addOutput(name)
645 return tens
646
Kevin Cheng550ccc52021-03-03 11:21:43 -0800647 def addOperator(self, op, inputs, outputs, attributes=None, quant_info=None):
Eric Kunzee5e26762020-10-13 16:11:07 -0700648
Kevin Cheng14d7f7a2021-05-12 10:44:49 -0700649 if op == tosa.Op.Op().CONST:
650 raise Exception("Use addConstTensor() to add CONST ops")
Eric Kunzee5e26762020-10-13 16:11:07 -0700651
Kevin Cheng550ccc52021-03-03 11:21:43 -0800652 return self.currBasicBlock.addOperator(
653 op, inputs, outputs, attributes, quant_info
654 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700655
Kevin Cheng550ccc52021-03-03 11:21:43 -0800656 def setExpectedFailure(self, desc="", val=True):
Eric Kunzee5e26762020-10-13 16:11:07 -0700657
Eric Kunzee5e26762020-10-13 16:11:07 -0700658 self.expectedFailure = val
659 self.expectedFailureDesc = desc
660
661 def serialize(self):
662
663 builder = self.builder
664
665 Version.VersionStart(builder)
666 Version.VersionAdd_major(builder, TOSA_VERSION[0])
667 Version.VersionAdd_minor(builder, TOSA_VERSION[1])
668 Version.VersionAdd_patch(builder, TOSA_VERSION[2])
669 Version.VersionAdd_experimental(builder, TOSA_VERSION[3])
670 version = Version.VersionEnd(builder)
671
Kevin Cheng550ccc52021-03-03 11:21:43 -0800672 fbv_bb = TosaSerializer.serializeObjVec(
673 builder, self.basicBlocks, TosaGraph.TosaGraphStartBlocksVector
674 )
Eric Kunzee5e26762020-10-13 16:11:07 -0700675
676 TosaGraph.TosaGraphStart(builder)
677 TosaGraph.TosaGraphAddVersion(builder, version)
678 TosaGraph.TosaGraphAddBlocks(builder, fbv_bb)
679 graph = TosaGraph.TosaGraphEnd(builder)
680
681 self.builder.Finish(graph)
682 return self.builder.Output()
683
684 def writeJson(self, tosa_filename):
Kevin Cheng550ccc52021-03-03 11:21:43 -0800685 """Write a json test file so that it is fairly easy to pick up the test
686 and generate commands for third party tool"""
Eric Kunzee5e26762020-10-13 16:11:07 -0700687 test_desc = dict()
688
Kevin Cheng550ccc52021-03-03 11:21:43 -0800689 test_desc["tosa_file"] = tosa_filename
Eric Kunzee5e26762020-10-13 16:11:07 -0700690 ifm_name = []
Eric Kunzee5e26762020-10-13 16:11:07 -0700691 ifm_file = []
692 ofm_name = []
693 ofm_file = []
Eric Kunzee5e26762020-10-13 16:11:07 -0700694
695 for b in self.basicBlocks:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800696 if b.name == "main":
Eric Kunzee5e26762020-10-13 16:11:07 -0700697 for i in b.inputs:
698 ifm_name.append(i)
Eric Kunzee5e26762020-10-13 16:11:07 -0700699 ifm_file.append(b.tensors[i].placeholderFilename)
700 for o in b.outputs:
701 ofm_name.append(o)
Eric Kunzee5e26762020-10-13 16:11:07 -0700702 # Make up an OFM filename here. One isn't generated until the reference tool is
703 # run, so any name is a good name
Kevin Cheng550ccc52021-03-03 11:21:43 -0800704 ofm_file.append("ref-{}.npy".format(o))
Eric Kunzee5e26762020-10-13 16:11:07 -0700705
Kevin Chengcd79f0e2021-06-03 15:00:34 -0700706 test_desc["ifm_name"] = ifm_name
Kevin Cheng550ccc52021-03-03 11:21:43 -0800707 test_desc["ifm_file"] = ifm_file
Kevin Cheng550ccc52021-03-03 11:21:43 -0800708 test_desc["ofm_name"] = ofm_name
Kevin Cheng550ccc52021-03-03 11:21:43 -0800709 test_desc["ofm_file"] = ofm_file
710 test_desc["expected_failure"] = self.expectedFailure
Eric Kunzee5e26762020-10-13 16:11:07 -0700711 if self.expectedFailureDesc:
Kevin Cheng550ccc52021-03-03 11:21:43 -0800712 test_desc["expected_failure_desc"] = self.expectedFailureDesc
Eric Kunzee5e26762020-10-13 16:11:07 -0700713
Kevin Cheng550ccc52021-03-03 11:21:43 -0800714 return json.dumps(test_desc, indent=" ")
Eric Kunzee5e26762020-10-13 16:11:07 -0700715
716 def startBasicBlock(self, name):
717 self.currBasicBlock = TosaSerializerBasicBlock(name)
718 self.basicBlocks.append(self.currBasicBlock)
719
720 @staticmethod
721 def serializeStrVec(builder, vec, start_fcn):
722 fb_strs = [builder.CreateString(i) for i in vec]
723 start_fcn(builder, len(fb_strs))
724 for s in fb_strs[::-1]:
725 builder.PrependUOffsetTRelative(s)
726 return builder.EndVector(len(fb_strs))
727
728 @staticmethod
Kevin Cheng82507d72021-06-17 16:01:59 -0700729 def serializeUint8Vec(builder, vec):
730 builder.StartVector(1, len(vec), 8)
731 for v in vec[::-1]:
732 builder.PrependUint8(v)
733 return builder.EndVector(len(vec))
734
735 @staticmethod
Eric Kunzee5e26762020-10-13 16:11:07 -0700736 def serializeInt32Vec(builder, vec):
737 builder.StartVector(4, len(vec), 4)
738 for v in vec[::-1]:
739 builder.PrependInt32(v)
740 return builder.EndVector(len(vec))
741
742 @staticmethod
Kevin Cheng77d0f762020-11-24 10:26:32 -0800743 def serializeFpVec(builder, vec):
744 builder.StartVector(4, len(vec), 4)
745 for v in vec[::-1]:
746 builder.PrependFloat32(v)
747 return builder.EndVector(len(vec))
748
749 @staticmethod
Eric Kunzee5e26762020-10-13 16:11:07 -0700750 def serializeObjVec(builder, vec, start_fcn):
751 serialized_vec = []
752 for v in vec[::-1]:
753 serialized_vec.append(v.serialize(builder))
754
755 start_fcn(builder, len(vec))
756 for v in serialized_vec:
757 builder.PrependUOffsetTRelative(v)
758 return builder.EndVector(len(vec))
759
760 @staticmethod
761 def toList(val):
762 if isinstance(val, list):
763 return val
764 else:
765 return [val]
766
767 @staticmethod
768 def setTosaVersion():
769 # Create a dummy flatbuffers file with the default version information
770 # There does not appear to be a better way to get a constant from a
771 # flatbuffer schema file
772 builder = flatbuffers.Builder(0)
773 Version.VersionStart(builder)
774 ver = Version.VersionEnd(builder)
775 TosaGraph.TosaGraphStart(builder)
776 TosaGraph.TosaGraphAddVersion(builder, ver)
777 gr = TosaGraph.TosaGraphEnd(builder)
778 builder.Finish(gr)
779
780 out = builder.Output()
781
782 gr = TosaGraph.TosaGraph()
783 root = gr.GetRootAsTosaGraph(out, 0)
784
785 # Store the version as a global variable so that it only needs to be
786 # generated once per process.
787 global TOSA_VERSION
Kevin Cheng550ccc52021-03-03 11:21:43 -0800788 TOSA_VERSION = [
789 root.Version()._major(),
790 root.Version()._minor(),
791 root.Version()._patch(),
792 root.Version()._experimental(),
793 ]