blob: fec676e63645aee768fb5883e621ab2daf1fb2e2 [file] [log] [blame]
Jeremy Johnson9b225172021-12-14 16:34:47 +00001# Copyright (c) 2020-2022, ARM Limited.
Kevin Chengfea5a372021-10-11 18:38:47 +00002#
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
Kevin Chengfea5a372021-10-11 18:38:47 +000015import os
Kevin Chengfea5a372021-10-11 18:38:47 +000016import json
17import flatbuffers
18import numpy as np
19import struct
Jeremy Johnson9b225172021-12-14 16:34:47 +000020from enum import IntEnum, unique
Kevin Chengfea5a372021-10-11 18:38:47 +000021from tosa import (
22 TosaGraph,
23 TosaBasicBlock,
24 TosaTensor,
25 TosaOperator,
Kevin Chengfea5a372021-10-11 18:38:47 +000026 Version,
27)
Jeremy Johnson9b225172021-12-14 16:34:47 +000028import tosa.DType as TosaDType
29import tosa.Op as TosaOp
Kevin Chengfea5a372021-10-11 18:38:47 +000030
Kevin Chenge6563f52021-10-20 12:12:02 -070031# Keep version number in sync with the version default value with schema/tosa.fbs
Kevin Chengb97cb1d2021-10-14 11:53:39 -070032TOSA_VERSION_MAJOR = 0
Eric Kunzea687b612021-11-03 17:02:57 -070033TOSA_VERSION_MINOR = 24
Kevin Chengb97cb1d2021-10-14 11:53:39 -070034TOSA_VERSION_PATCH = 0
Eric Kunzea687b612021-11-03 17:02:57 -070035TOSA_VERSION_DRAFT = True
Jeremy Johnson9b225172021-12-14 16:34:47 +000036TOSA_VERSION = [
37 TOSA_VERSION_MAJOR,
38 TOSA_VERSION_MINOR,
39 TOSA_VERSION_PATCH,
40 TOSA_VERSION_DRAFT,
41]
Kevin Chengfea5a372021-10-11 18:38:47 +000042# 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.
Jeremy Johnson9b225172021-12-14 16:34:47 +000045DType = TosaDType.DType()
Kevin Chengfea5a372021-10-11 18:38:47 +000046DTypeNames = [
47 "UNKNOWN",
48 "BOOL",
49 "UINT8",
50 "INT4",
51 "INT8",
52 "INT16",
53 "INT32",
54 "INT48",
55 "FLOAT",
Jeremy Johnson41027732022-05-25 17:52:29 +010056 "UINT16",
Kevin Chengfea5a372021-10-11 18:38:47 +000057]
58
59ByteMask = np.uint64(0xFF)
60
61
62def dtype_str_to_val(name):
63
64 for i in range(len(DTypeNames)):
65 if name.casefold() == DTypeNames[i].casefold():
66 return i
67 raise Exception("Unable to parse DType name {}".format(name))
68
69
70class TosaSerializerUnion:
71 """This class handles encapsulating and serializing union types into flatbuffers"""
72
73 def __init__(self):
74
Jeremy Johnson9b225172021-12-14 16:34:47 +000075 # A tuple of the start and end functions.
76 # Set by the options constructors below
Kevin Chengfea5a372021-10-11 18:38:47 +000077 self.optFcns = None
78
Jeremy Johnson9b225172021-12-14 16:34:47 +000079 # The type from the tosa.Options enumeration.
80 # Set by the options constructors below.
Kevin Chengfea5a372021-10-11 18:38:47 +000081 self.utype = None
82
83 # Each of these lists is a tuple of the add function and the
84 # value being added. Set by the options constructors below.
85 self.ints = []
86 self.bools = []
87 self.floats = []
88 self.strings = []
89 self.intvecs = []
90 self.fpvecs = []
91
92 def serialize(self, builder):
93
94 # We have to build strings and vectors first
95 strList = []
96 intVecList = []
97 fpVecList = []
98
99 for fcn, val in self.strings:
100 strList.append((fcn, builder.CreateString(val)))
101
102 for fcn, val in self.intvecs:
103 intVecList.append((fcn, TosaSerializer.serializeInt32Vec(builder, val)))
104
105 for fcn, val in self.fpvecs:
106 fpVecList.append((fcn, TosaSerializer.serializeFpVec(builder, val)))
107
108 startFcn, endFcn = self.optFcns
109
110 # Then serialize the options object from the list of primitives and
111 # other serialized values
112 startFcn(builder)
113 for fcn, val in self.ints:
114 fcn(builder, val)
115
116 for fcn, val in self.bools:
117 fcn(builder, val)
118
119 for fcn, val in self.floats:
120 fcn(builder, val)
121
122 for fcn, val in strList:
123 fcn(builder, val)
124
125 for fcn, val in intVecList:
126 fcn(builder, val)
127
128 for fcn, val in fpVecList:
129 fcn(builder, val)
130
131 return endFcn(builder)
132
133
134class TosaSerializerAttribute(TosaSerializerUnion):
135 """This class handles encapsulating all of the enumerated types for attributes"""
136
137 def __init__(self):
138 super().__init__()
139
TatWai Chong7be71652022-05-10 17:26:20 -0700140 def PoolAttribute(self, kernel, stride, pad):
Kevin Chengfea5a372021-10-11 18:38:47 +0000141 from tosa import PoolAttribute as a, Attribute
142
143 self.utype = Attribute.Attribute().PoolAttribute
144
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800145 self.optFcns = (a.Start, a.End)
TatWai Chong7be71652022-05-10 17:26:20 -0700146 self.intvecs.append((a.AddPad, pad))
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800147 self.intvecs.append((a.AddKernel, kernel))
148 self.intvecs.append((a.AddStride, stride))
Kevin Chengfea5a372021-10-11 18:38:47 +0000149
TatWai Chong7be71652022-05-10 17:26:20 -0700150 def ConvAttribute(self, pad, stride, dilation):
Kevin Chengfea5a372021-10-11 18:38:47 +0000151 from tosa import ConvAttribute as a, Attribute
152
153 self.utype = Attribute.Attribute().ConvAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800154 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000155
TatWai Chong7be71652022-05-10 17:26:20 -0700156 self.intvecs.append((a.AddPad, pad))
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800157 self.intvecs.append((a.AddStride, stride))
158 self.intvecs.append((a.AddDilation, dilation))
Kevin Chengfea5a372021-10-11 18:38:47 +0000159
160 def TransposeConvAttribute(self, outpad, stride, dilation, output_shape):
161 from tosa import TransposeConvAttribute as a, Attribute
162
163 self.utype = Attribute.Attribute().TransposeConvAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800164 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000165
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800166 self.intvecs.append((a.AddOutpad, outpad))
167 self.intvecs.append((a.AddStride, stride))
168 self.intvecs.append((a.AddDilation, dilation))
169 self.intvecs.append((a.AddOutputShape, output_shape))
Kevin Chengfea5a372021-10-11 18:38:47 +0000170
Kevin Cheng38d214c2021-10-15 15:49:19 -0700171 def PadAttribute(self, padding, pad_const_int, pad_const_fp):
172 from tosa import PadAttribute as a, Attribute
Kevin Chengfea5a372021-10-11 18:38:47 +0000173
Kevin Cheng38d214c2021-10-15 15:49:19 -0700174 self.utype = Attribute.Attribute().PadAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800175 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000176
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800177 self.intvecs.append((a.AddPadding, padding))
178 self.ints.append((a.AddPadConstInt, pad_const_int))
179 self.floats.append((a.AddPadConstFp, pad_const_fp))
Kevin Chengfea5a372021-10-11 18:38:47 +0000180
181 def AxisAttribute(self, axis):
182 from tosa import AxisAttribute as a, Attribute
183
184 self.utype = Attribute.Attribute().AxisAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800185 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000186
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800187 self.ints.append((a.AddAxis, axis))
Kevin Chengfea5a372021-10-11 18:38:47 +0000188
TatWai Chong7be71652022-05-10 17:26:20 -0700189 def ReshapeAttribute(self, new_shape):
Kevin Chengfea5a372021-10-11 18:38:47 +0000190 from tosa import ReshapeAttribute as a, Attribute
191
192 self.utype = Attribute.Attribute().ReshapeAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800193 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000194
TatWai Chong7be71652022-05-10 17:26:20 -0700195 self.intvecs.append((a.AddNewShape, new_shape))
Kevin Chengfea5a372021-10-11 18:38:47 +0000196
TatWai Chong7be71652022-05-10 17:26:20 -0700197 def SliceAttribute(self, start, size):
Kevin Chengfea5a372021-10-11 18:38:47 +0000198 from tosa import SliceAttribute as a, Attribute
199
200 self.utype = Attribute.Attribute().SliceAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800201 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000202
TatWai Chong7be71652022-05-10 17:26:20 -0700203 self.intvecs.append((a.AddStart, start))
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800204 self.intvecs.append((a.AddSize, size))
Kevin Chengfea5a372021-10-11 18:38:47 +0000205
206 def TileAttribute(self, multiples):
207 from tosa import TileAttribute as a, Attribute
208
209 self.utype = Attribute.Attribute().TileAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800210 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000211
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800212 self.intvecs.append((a.AddMultiples, multiples))
Kevin Chengfea5a372021-10-11 18:38:47 +0000213
214 def ResizeAttribute(
215 self, output_size, stride, offset, shift, stride_fp, offset_fp, mode
216 ):
217 from tosa import ResizeAttribute as a, Attribute
218
219 self.utype = Attribute.Attribute().ResizeAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800220 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000221
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800222 self.intvecs.append((a.AddOutputSize, output_size))
223 self.intvecs.append((a.AddStride, stride))
224 self.intvecs.append((a.AddOffset, offset))
225 self.ints.append((a.AddShift, shift))
226 self.fpvecs.append((a.AddStrideFp, stride_fp))
227 self.fpvecs.append((a.AddOffsetFp, offset_fp))
228 self.ints.append((a.AddMode, mode))
Kevin Chengfea5a372021-10-11 18:38:47 +0000229
230 def ClampAttribute(self, minint, maxint, minfp, maxfp):
231 from tosa import ClampAttribute as a, Attribute
232
233 self.utype = Attribute.Attribute().ClampAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800234 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000235
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800236 self.ints.append((a.AddMinInt, minint))
237 self.ints.append((a.AddMaxInt, maxint))
Kevin Chengfea5a372021-10-11 18:38:47 +0000238
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800239 self.ints.append((a.AddMinFp, minfp))
240 self.ints.append((a.AddMaxFp, maxfp))
Kevin Chengfea5a372021-10-11 18:38:47 +0000241
242 def RescaleAttribute(
243 self, input_zp, output_zp, multiplier, shift, scale32, double_round, per_channel
244 ):
245 from tosa import RescaleAttribute as a, Attribute
246
247 self.utype = Attribute.Attribute().RescaleAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800248 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000249
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800250 self.ints.append((a.AddInputZp, input_zp))
251 self.ints.append((a.AddOutputZp, output_zp))
252 self.intvecs.append((a.AddMultiplier, multiplier))
253 self.intvecs.append((a.AddShift, shift))
254 self.bools.append((a.AddScale32, scale32))
255 self.bools.append((a.AddDoubleRound, double_round))
256 self.bools.append((a.AddPerChannel, per_channel))
Kevin Chengfea5a372021-10-11 18:38:47 +0000257
258 def MulAttribute(self, shift):
259 from tosa import MulAttribute as a, Attribute
260
261 self.utype = Attribute.Attribute().MulAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800262 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000263
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800264 self.ints.append((a.AddShift, shift))
Kevin Chengfea5a372021-10-11 18:38:47 +0000265
266 def ArithmeticRightShiftAttribute(self, round):
267 from tosa import ArithmeticRightShiftAttribute as a, Attribute
268
269 self.utype = Attribute.Attribute().ArithmeticRightShiftAttribute
270 self.optFcns = (
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800271 a.Start,
272 a.End,
Kevin Chengfea5a372021-10-11 18:38:47 +0000273 )
274
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800275 self.bools.append((a.AddRound, round))
Kevin Chengfea5a372021-10-11 18:38:47 +0000276
Kevin Chengfea5a372021-10-11 18:38:47 +0000277 def CondIfAttribute(self, then_branch, else_branch):
278 from tosa import CondIfAttribute as a, Attribute
279
280 self.utype = Attribute.Attribute().CondIfAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800281 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000282
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800283 self.strings.append((a.AddThenBranch, then_branch))
284 self.strings.append((a.AddElseBranch, else_branch))
Kevin Chengfea5a372021-10-11 18:38:47 +0000285
286 def WhileLoopAttribute(self, cond_branch, body_branch):
287 from tosa import WhileLoopAttribute as a, Attribute
288
289 self.utype = Attribute.Attribute().WhileLoopAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800290 self.optFcns = (a.Start, a.End)
Kevin Chengfea5a372021-10-11 18:38:47 +0000291
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800292 self.strings.append((a.AddCondBranch, cond_branch))
293 self.strings.append((a.AddBodyBranch, body_branch))
Kevin Chengfea5a372021-10-11 18:38:47 +0000294
TatWai Chong7be71652022-05-10 17:26:20 -0700295 def TransposeAttribute(self, perms):
Kevin Cheng38d214c2021-10-15 15:49:19 -0700296 from tosa import TransposeAttribute as a, Attribute
297
298 self.utype = Attribute.Attribute().TransposeAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800299 self.optFcns = (a.Start, a.End)
Kevin Cheng38d214c2021-10-15 15:49:19 -0700300
TatWai Chong7be71652022-05-10 17:26:20 -0700301 self.intvecs.append((a.AddPerms, perms))
Kevin Cheng38d214c2021-10-15 15:49:19 -0700302
303 def TableAttribute(self, table):
304 from tosa import TableAttribute as a, Attribute
305
306 self.utype = Attribute.Attribute().TableAttribute
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800307 self.optFcns = (a.Start, a.End)
Kevin Cheng38d214c2021-10-15 15:49:19 -0700308
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800309 self.intvecs.append((a.AddTable, table))
Kevin Chengfea5a372021-10-11 18:38:47 +0000310
Jeremy Johnson9b225172021-12-14 16:34:47 +0000311
Kevin Chengfea5a372021-10-11 18:38:47 +0000312class TosaSerializerQuantInfo(TosaSerializerUnion):
Jeremy Johnson9b225172021-12-14 16:34:47 +0000313 """This class handles encapsulating all of the enumerated types for quantinfo"""
Kevin Chengfea5a372021-10-11 18:38:47 +0000314
315 def __init__(self):
316 super().__init__()
317
318 def ConvQuantInfo(self, input_zp, weight_zp):
319 from tosa import ConvQuantInfo as q, QuantInfo
320
321 self.utype = QuantInfo.QuantInfo().ConvQuantInfo
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800322 self.optFcns = (q.Start, q.End)
323 self.ints.append((q.AddInputZp, input_zp))
324 self.ints.append((q.AddWeightZp, weight_zp))
Kevin Chengfea5a372021-10-11 18:38:47 +0000325
326 def UnaryQuantInfo(self, input_zp, output_zp):
327 from tosa import UnaryQuantInfo as q, QuantInfo
328
329 self.utype = QuantInfo.QuantInfo().UnaryQuantInfo
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800330 self.optFcns = (q.Start, q.End)
331 self.ints.append((q.AddInputZp, input_zp))
332 self.ints.append((q.AddOutputZp, output_zp))
Kevin Chengfea5a372021-10-11 18:38:47 +0000333
334 def MatMulQuantInfo(self, a_zp, b_zp):
335 from tosa import MatMulQuantInfo as q, QuantInfo
336
337 self.utype = QuantInfo.QuantInfo().MatMulQuantInfo
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800338 self.optFcns = (q.Start, q.End)
339 self.ints.append((q.AddAZp, a_zp))
340 self.ints.append((q.AddBZp, b_zp))
Kevin Chengfea5a372021-10-11 18:38:47 +0000341
342 def PadQuantInfo(self, input_zp):
343 from tosa import PadQuantInfo as q, QuantInfo
344
345 self.utype = QuantInfo.QuantInfo().PadQuantInfo
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800346 self.optFcns = (q.Start, q.End)
347 self.ints.append((q.AddInputZp, input_zp))
Kevin Chengfea5a372021-10-11 18:38:47 +0000348
349
350class TosaSerializerTensor:
351 def __init__(
352 self,
353 name,
354 shape,
355 dtype,
356 data=None,
357 placeholderFilename=None,
358 ):
359 self.name = name
360
361 if isinstance(shape, np.ndarray):
362 shape = shape.astype(int).tolist()
363 shape = list(map(int, shape))
364
365 self.shape = shape
366 self.dtype = dtype
367
368 if isinstance(data, np.ndarray):
369 data = data.flatten().astype(int).tolist()
370 data = list(map(int, data))
371 self.data = data
372 elif isinstance(data, list):
373 data = list(map(int, data))
374 self.data = data
375 else:
376 self.data = None
377
378 # Filename for placeholder tensors. These get generated by the test generation
Jeremy Johnson9b225172021-12-14 16:34:47 +0000379 # process and are written to disk, but are considered input tensors by the
380 # network so they do not appear in the TOSA serialiazation. However, if we
381 # want to form a unit test around these input tensors, we can get the filename
382 # from here.
Kevin Chengfea5a372021-10-11 18:38:47 +0000383 self.placeholderFilename = placeholderFilename
384
385 def __str__(self):
386 str = "TosaSerializerTensor name: {} shape: {} dtype: {}".format(
387 self.name,
388 self.shape,
389 DTypeNames[self.dtype],
390 )
391 return str
392
393 def setDtype(self, dtype):
394 self.dtype = dtype
395
396 def serialize(self, builder):
397 fb_name = builder.CreateString(self.name)
398 fb_shapes = TosaSerializer.serializeInt32Vec(builder, self.shape)
399 if self.data:
400 u8_data = list()
401 # little endianess
402 if self.dtype == DType.BOOL:
403 for val in self.data:
404 val_u8 = np.uint8(val)
405 u8_data.append(val_u8)
406 elif self.dtype == DType.INT4:
407 in_size = len(self.data)
408 out_size = (in_size + 1) // 2
409 for i in range(out_size):
410 val_0 = self.data[2 * i]
411 if (2 * i + 1) < in_size:
412 val_1 = self.data[2 * i + 1]
413 else:
414 val_1 = 0
415 val_i8 = (val_0 & 0xF) | ((val_1 & 0xF) << 4)
416 val_u8 = np.uint8(val_i8)
417 u8_data.append(val_u8)
418 elif self.dtype == DType.INT8:
419 for val in self.data:
420 val_u8 = np.uint8(val)
421 u8_data.append(val_u8)
422 elif self.dtype == DType.INT16:
423 for val in self.data:
424 val_u16 = np.uint16(val)
425 b0 = val_u16 & ByteMask
426 b1 = (val_u16 >> np.uint16(8)) & ByteMask
427 u8_data.extend([b0, b1])
428 elif self.dtype == DType.INT32:
429 for val in self.data:
430 val_u32 = np.uint32(val)
431 b0 = val_u32 & ByteMask
432 b1 = (val_u32 >> np.uint32(8)) & ByteMask
433 b2 = (val_u32 >> np.uint32(16)) & ByteMask
Kevin Cheng6b078ca2021-10-13 23:12:50 -0700434 b3 = (val_u32 >> np.uint32(24)) & ByteMask
Kevin Chengfea5a372021-10-11 18:38:47 +0000435 u8_data.extend([b0, b1, b2, b3])
436 elif self.dtype == DType.INT48:
437 for val in self.data:
438 val_u64 = np.uint64(val)
439 b0 = val_u64 & ByteMask
440 b1 = (val_u64 >> np.uint64(8)) & ByteMask
441 b2 = (val_u64 >> np.uint64(16)) & ByteMask
442 b3 = (val_u64 >> np.uint64(24)) & ByteMask
443 b4 = (val_u64 >> np.uint64(32)) & ByteMask
444 b5 = (val_u64 >> np.uint64(40)) & ByteMask
445 u8_data.extend([b0, b1, b2, b3, b4, b5])
446 elif self.dtype == DType.FLOAT:
447 for val in self.data:
448 b = struct.pack("!f", val)
449 u8_data.extend([b[3], b[2], b[1], b[0]])
450 else:
451 raise Exception(
452 "unsupported data type {}".format(DTypeNames[self.dtype])
453 )
454 fb_data = TosaSerializer.serializeUint8Vec(builder, u8_data)
455
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800456 TosaTensor.Start(builder)
457 TosaTensor.AddName(builder, fb_name)
458 TosaTensor.AddShape(builder, fb_shapes)
459 TosaTensor.AddType(builder, self.dtype)
Kevin Chengfea5a372021-10-11 18:38:47 +0000460 if self.data:
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800461 TosaTensor.AddData(builder, fb_data)
Kevin Chengfea5a372021-10-11 18:38:47 +0000462
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800463 return TosaTensor.End(builder)
Kevin Chengfea5a372021-10-11 18:38:47 +0000464
465
466class TosaSerializerOperator:
467 def __init__(self, op, inputs, outputs, attributes=None, quantInfo=None):
468 self.op = op
469 self.attributes = attributes
470 self.inputs = TosaSerializer.toList(inputs)
471 self.outputs = TosaSerializer.toList(outputs)
472 self.quantInfo = quantInfo
473
474 def __str__(self):
475 str = "Op {}\n----\n".format(self.op)
476
477 for i in self.inputs:
478 str = str + " Input: {}\n".format(i)
479 for o in self.outputs:
480 str = str + " Output: {}\n".format(o)
481
482 return str
483
484 def serialize(self, builder):
485 fb_inputs = TosaSerializer.serializeStrVec(
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800486 builder, self.inputs, TosaOperator.StartInputsVector
Kevin Chengfea5a372021-10-11 18:38:47 +0000487 )
488 fb_outputs = TosaSerializer.serializeStrVec(
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800489 builder, self.outputs, TosaOperator.StartOutputsVector
Kevin Chengfea5a372021-10-11 18:38:47 +0000490 )
491 # Need to serialize quant_info and attributes enums still
492 if self.attributes is not None:
493 fb_attributes = self.attributes.serialize(builder)
494
495 if self.quantInfo is not None:
496 fb_qinfo = self.quantInfo.serialize(builder)
497
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800498 TosaOperator.Start(builder)
499 TosaOperator.AddOp(builder, self.op)
500 TosaOperator.AddInputs(builder, fb_inputs)
501 TosaOperator.AddOutputs(builder, fb_outputs)
Kevin Chengfea5a372021-10-11 18:38:47 +0000502 if self.attributes is not None:
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800503 TosaOperator.AddAttributeType(builder, self.attributes.utype)
504 TosaOperator.AddAttribute(builder, fb_attributes)
Kevin Chengfea5a372021-10-11 18:38:47 +0000505 if self.quantInfo is not None:
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800506 TosaOperator.AddQuantInfoType(builder, self.quantInfo.utype)
507 TosaOperator.AddQuantInfo(builder, fb_qinfo)
Kevin Chengfea5a372021-10-11 18:38:47 +0000508
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800509 return TosaOperator.End(builder)
Kevin Chengfea5a372021-10-11 18:38:47 +0000510
511
512class TosaSerializerBasicBlock:
513 def __init__(self, name):
514 self.name = name
515 self.operators = []
516
517 # Dict assures uniqueness, but allows us to look up by name
518 self.tensors = dict()
519
520 self.inputs = []
521 self.outputs = []
522
523 def addTensor(
524 self,
525 name,
526 shape,
527 dtype,
528 data=None,
529 placeholderFilename=None,
530 ):
Jeremy Johnson9b225172021-12-14 16:34:47 +0000531 if name not in self.tensors:
Kevin Chengfea5a372021-10-11 18:38:47 +0000532 self.tensors[name] = TosaSerializerTensor(
533 name, shape, dtype, data, placeholderFilename
534 )
535
536 return self.tensors[name]
537
538 def addInput(self, name):
539 self.inputs.append(name)
540
541 def addOutput(self, name):
542 self.outputs.append(name)
543
544 def addOperator(self, op, inputs, outputs, attributes=None, quant_info=None):
545 self.operators.append(
546 TosaSerializerOperator(op, inputs, outputs, attributes, quant_info)
547 )
548
549 def serialize(self, builder):
550 fb_name = builder.CreateString(self.name)
551 fbv_inputs = TosaSerializer.serializeStrVec(
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800552 builder, list(self.inputs), TosaBasicBlock.StartInputsVector
Kevin Chengfea5a372021-10-11 18:38:47 +0000553 )
554 fbv_outputs = TosaSerializer.serializeStrVec(
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800555 builder, list(self.outputs), TosaBasicBlock.StartOutputsVector
Kevin Chengfea5a372021-10-11 18:38:47 +0000556 )
557 fbv_tensors = TosaSerializer.serializeObjVec(
558 builder,
559 list(self.tensors.values()),
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800560 TosaBasicBlock.StartTensorsVector,
Kevin Chengfea5a372021-10-11 18:38:47 +0000561 )
562 fbv_operators = TosaSerializer.serializeObjVec(
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800563 builder, self.operators, TosaBasicBlock.StartOperatorsVector
Kevin Chengfea5a372021-10-11 18:38:47 +0000564 )
565
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800566 TosaBasicBlock.Start(builder)
567 TosaBasicBlock.AddName(builder, fb_name)
568 TosaBasicBlock.AddInputs(builder, fbv_inputs)
569 TosaBasicBlock.AddOutputs(builder, fbv_outputs)
570 TosaBasicBlock.AddTensors(builder, fbv_tensors)
571 TosaBasicBlock.AddOperators(builder, fbv_operators)
572 return TosaBasicBlock.End(builder)
Kevin Chengfea5a372021-10-11 18:38:47 +0000573
574
575@unique
576class TensorDir(IntEnum):
577 PLACEHOLDER = 0
578 CONST = 1
579 INTERMEDIATE = 2
580 RESULT = 3
581
582
583class TosaSerializer:
584 def __init__(self, pathPrefix):
585
586 # Get the global TOSA version if not already defined
Kevin Chengfea5a372021-10-11 18:38:47 +0000587
588 self.builder = flatbuffers.Builder(0)
589
590 self.basicBlocks = []
591 self.startBasicBlock("main")
592 self.pathPrefix = pathPrefix
593
594 # Indicies used for adding/naming tensors
595 self.currInputIdx = 0
596 self.currConstIdx = 0
597 self.currLayerIdx = 1
598 self.currResultIdx = 0
599
600 # Is this an illegal test that is expected to fail?
Jeremy Johnson9b225172021-12-14 16:34:47 +0000601 self.expectedReturnCode = 0
Kevin Chengfea5a372021-10-11 18:38:47 +0000602 self.expectedFailure = False
603 self.expectedFailureDesc = ""
604
605 def __str__(self):
606 str = ""
607 for bb in self.basicBlocks:
608 str = str + bb.__str__()
609 return str
610
611 def addPlaceholder(self, shape, dtype, vals):
612 if not self.currBasicBlock:
613 raise Exception("addTensor called without valid basic block")
614
615 name = "input-{}".format(self.currInputIdx)
616 filename = "{}.npy".format(name)
617 self.currInputIdx = self.currInputIdx + 1
618
619 tens = self.currBasicBlock.addTensor(name, shape, dtype, None, filename)
620 # This is always an input to the block
621 self.currBasicBlock.addInput(name)
622
623 if vals is not None:
624 np.save(os.path.join(self.pathPrefix, filename), vals, False)
625
626 return tens
627
628 def addConst(self, shape, dtype, vals):
629 if not self.currBasicBlock:
630 raise Exception("addTensor called without valid basic block")
631
632 name = "const-{}".format(self.currInputIdx)
Kevin Chengfea5a372021-10-11 18:38:47 +0000633 self.currInputIdx = self.currInputIdx + 1
634
635 tens = self.currBasicBlock.addTensor(name, shape, dtype, vals)
636 # Add the operator now
Jeremy Johnson9b225172021-12-14 16:34:47 +0000637 self.currBasicBlock.addOperator(TosaOp.Op().CONST, [], name)
Kevin Chengfea5a372021-10-11 18:38:47 +0000638
639 return tens
640
641 def addIntermediate(self, shape, dtype):
642
643 if not self.currBasicBlock:
644 raise Exception("addTensor called without valid basic block")
645
646 name = "layer-{}".format(self.currLayerIdx)
647 self.currLayerIdx = self.currLayerIdx + 1
648
649 tens = self.currBasicBlock.addTensor(name, shape, dtype, None)
650
651 return tens
652
653 def addInputTensor(self, tensor):
654 self.currBasicBlock.addTensor(tensor.name, tensor.shape, tensor.dtype)
655 self.currBasicBlock.addInput(tensor.name)
656
657 def addOutputTensor(self, tensor):
658 self.currBasicBlock.addOutput(tensor.name)
659
660 def addOutput(self, shape, dtype):
661 if not self.currBasicBlock:
662 raise Exception("addTensor called without valid basic block")
663
664 name = "result-{}".format(self.currResultIdx)
665 self.currResultIdx = self.currResultIdx + 1
666
667 tens = self.currBasicBlock.addTensor(name, shape, dtype, None)
668 self.currBasicBlock.addOutput(name)
669 return tens
670
671 def addOperator(self, op, inputs, outputs, attributes=None, quant_info=None):
672
Jeremy Johnson9b225172021-12-14 16:34:47 +0000673 if op == TosaOp.Op().CONST:
Kevin Chengfea5a372021-10-11 18:38:47 +0000674 raise Exception("Use addConstTensor() to add CONST ops")
675
676 return self.currBasicBlock.addOperator(
677 op, inputs, outputs, attributes, quant_info
678 )
679
Jeremy Johnson9b225172021-12-14 16:34:47 +0000680 def setExpectedReturnCode(self, val, fail, desc=""):
Kevin Chengfea5a372021-10-11 18:38:47 +0000681
682 self.expectedReturnCode = val
683 self.expectedFailureDesc = desc
Jeremy Johnson9b225172021-12-14 16:34:47 +0000684 self.expectedFailure = fail
Kevin Chengfea5a372021-10-11 18:38:47 +0000685
686 def serialize(self):
687
688 builder = self.builder
689
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800690 Version.Start(builder)
691 Version.Add_major(builder, TOSA_VERSION[0])
692 Version.Add_minor(builder, TOSA_VERSION[1])
693 Version.Add_patch(builder, TOSA_VERSION[2])
694 Version.Add_draft(builder, TOSA_VERSION[3])
695 version = Version.End(builder)
Kevin Chengfea5a372021-10-11 18:38:47 +0000696
697 fbv_bb = TosaSerializer.serializeObjVec(
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800698 builder, self.basicBlocks, TosaGraph.StartBlocksVector
Kevin Chengfea5a372021-10-11 18:38:47 +0000699 )
700
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800701 TosaGraph.Start(builder)
702 TosaGraph.AddVersion(builder, version)
703 TosaGraph.AddBlocks(builder, fbv_bb)
704 graph = TosaGraph.End(builder)
Kevin Chengfea5a372021-10-11 18:38:47 +0000705
706 self.builder.Finish(graph)
707 return self.builder.Output()
708
709 def writeJson(self, tosa_filename):
710 """Write a json test file so that it is fairly easy to pick up the test
711 and generate commands for third party tool"""
712 test_desc = dict()
713
714 test_desc["tosa_file"] = tosa_filename
715 ifm_name = []
716 ifm_file = []
717 ofm_name = []
718 ofm_file = []
719
720 for b in self.basicBlocks:
721 if b.name == "main":
722 for i in b.inputs:
723 ifm_name.append(i)
724 ifm_file.append(b.tensors[i].placeholderFilename)
725 for o in b.outputs:
726 ofm_name.append(o)
Jeremy Johnson9b225172021-12-14 16:34:47 +0000727 # Make up an OFM filename here. One isn't generated until the
728 # reference tool is run, so any name is a good name
Kevin Chengfea5a372021-10-11 18:38:47 +0000729 ofm_file.append("ref-{}.npy".format(o))
730
731 test_desc["ifm_name"] = ifm_name
732 test_desc["ifm_file"] = ifm_file
733 test_desc["ofm_name"] = ofm_name
734 test_desc["ofm_file"] = ofm_file
735 test_desc["expected_return_code"] = self.expectedReturnCode
736 test_desc["expected_failure"] = self.expectedFailure
737 if self.expectedFailureDesc:
738 test_desc["expected_failure_desc"] = self.expectedFailureDesc
739
740 return json.dumps(test_desc, indent=" ")
741
742 def startBasicBlock(self, name):
743 self.currBasicBlock = TosaSerializerBasicBlock(name)
744 self.basicBlocks.append(self.currBasicBlock)
745
746 @staticmethod
747 def serializeStrVec(builder, vec, start_fcn):
748 fb_strs = [builder.CreateString(i) for i in vec]
749 start_fcn(builder, len(fb_strs))
750 for s in fb_strs[::-1]:
751 builder.PrependUOffsetTRelative(s)
Kevin Cheng49faa4e2021-11-08 16:59:18 -0800752 return builder.EndVector()
Kevin Chengfea5a372021-10-11 18:38:47 +0000753
754 @staticmethod
755 def serializeUint8Vec(builder, vec):
756 builder.StartVector(1, len(vec), 8)
757 for v in vec[::-1]:
758 builder.PrependUint8(v)
759 try:
760 return builder.EndVector()
761 except TypeError:
762 return builder.EndVector(len(vec))
763
764 @staticmethod
765 def serializeInt32Vec(builder, vec):
766 builder.StartVector(4, len(vec), 4)
767 for v in vec[::-1]:
768 builder.PrependInt32(v)
769 try:
770 return builder.EndVector()
771 except TypeError:
772 return builder.EndVector(len(vec))
773
774 @staticmethod
775 def serializeFpVec(builder, vec):
776 builder.StartVector(4, len(vec), 4)
777 for v in vec[::-1]:
778 builder.PrependFloat32(v)
779 try:
780 return builder.EndVector()
781 except TypeError:
782 return builder.EndVector(len(vec))
783
784 @staticmethod
785 def serializeObjVec(builder, vec, start_fcn):
786 serialized_vec = []
787 for v in vec[::-1]:
788 serialized_vec.append(v.serialize(builder))
789
790 start_fcn(builder, len(vec))
791 for v in serialized_vec:
792 builder.PrependUOffsetTRelative(v)
793 try:
794 return builder.EndVector()
795 except TypeError:
796 return builder.EndVector(len(vec))
797
798 @staticmethod
799 def toList(val):
800 if isinstance(val, list):
801 return val
802 else:
803 return [val]