#########################################################
# Copyright (C) 2022 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Christopher Rodrigues
#########################################################
"""
Codec implementations for Awesome IR data types.
Generally, whenever there is a change to the data stored
in an Awesome IR data type, the codec for that type should
change to properly load and store the data.
"""
from typing import Dict, Any, Sequence, Generic, TypeVar, List, Tuple, Union, Optional
import numpy as np
import json
import ml_kernels.requantization as requantization
from afe.core.graph_analyzer.utils import Metric
from ml_kernels.math_helpers import RoundType, bfloat16
from afe import backends
from afe.apis.defines import (
ChromaSampling, ColorConversion, ColorSpaceStandard, ResizeDepositLocation, ResizeMethod
)
from afe.backends import Backend
from afe.common_utils import is_installed
from afe.core.configs import QuantizationConfigs, QuantizationPrecision, Opt, EmptyValue
from afe.ir import operations
from afe.ir import attributes
from afe.ir.attributes import QuantResultTensorType
from afe.ir.defines import Status, NodeName, Quantization, DataValue, TensorValue, TupleValue, NodeAssociatedValue, \
LayerStats, RequantMethod, RequantizationMode, BiasCorrectionType
from afe.ir.net import AwesomeNet, AwesomeNode
from afe.ir.serializer.yaml_codec import (
map_codec, named_dict_codec, scalar_codec, list_codec, homogeneous_tuple_codec, singleton_dict_codec,
intrusive_tagged_union_codec, enum_codec, type_tagged_codec, dataclass_style_codec, FieldEncoding, Codec,
EncStateV, DecStateV, optional_codec, dict_codec, lazy_codec, DecodeError, adt_style_codec, ConstructorEncoding,
EncodeError, EncState, DecState, SCALAR_TYPES, null_value_codec
)
from afe.ir.sima_ir import SiMaIRMetadata, SiMaIR
from afe.ir.tensor_type import TensorType, ScalarType
from sima_utils.logging import sima_logger
from sima_utils.common import Platform
_Scalar = TypeVar("_Scalar", int, bool, float, str, None)
_A = TypeVar("_A")
_B = TypeVar("_B")
# Codec[EncStateV, DecStateV, List[_Scalar]]
_list_scalar_codec = list_codec(scalar_codec)
# Codec[EncStateV, DecStateV, Tuple[_Scalar, ...]]
_tuple_scalar_codec = homogeneous_tuple_codec(scalar_codec)
_list_tuple_scalar_codec = list_codec(_tuple_scalar_codec)
class _ScalarOrListScalarCodec(Generic[EncStateV, DecStateV, _Scalar]):
def encode(self, state: EncStateV, obj: Union[_Scalar, List[_Scalar]]) -> Any:
if isinstance(obj, list):
return _list_scalar_codec.encode(state, obj)
else:
return scalar_codec.encode(state, obj)
def decode(self, state: DecStateV, doc: Any) -> Union[_Scalar, List[_Scalar]]:
if isinstance(doc, list):
return _list_scalar_codec.decode(state, doc)
else:
return scalar_codec.decode(state, doc)
_scalar_or_list_scalar_codec = _ScalarOrListScalarCodec()
class _ScalarOrTupleScalarCodec(Generic[EncStateV, DecStateV, _Scalar]):
def encode(self, state: EncStateV, obj: Union[_Scalar, Tuple[_Scalar, ...]]) -> Any:
if isinstance(obj, tuple):
return _tuple_scalar_codec.encode(state, obj)
else:
return scalar_codec.encode(state, obj)
def decode(self, state: DecStateV, doc: Any) -> Union[_Scalar, Tuple[_Scalar, ...]]:
if isinstance(doc, list): # Tuple is encoded as list
return _tuple_scalar_codec.decode(state, doc)
else:
return scalar_codec.decode(state, doc)
_scalar_or_tuple_scalar_codec = _ScalarOrTupleScalarCodec()
class _ScalarOrNDArrayCodec(Generic[EncStateV, DecStateV]):
def encode(self, state: EncStateV, obj: Union[_Scalar, np.ndarray]) -> Any:
if isinstance(obj, np.ndarray):
return ndarray_codec.encode(state, obj)
else:
return scalar_codec.encode(state, obj)
def decode(self, state: DecStateV, doc: Any) -> Union[_Scalar, np.ndarray]:
if isinstance(doc, str): # Array is encoded as string
return ndarray_codec.decode(state, doc)
else:
return scalar_codec.decode(state, doc)
_scalar_or_ndarray_codec = _ScalarOrNDArrayCodec()
[docs]
ROUNDTYPE_CODEC: Codec[EncState, DecState, RoundType] = enum_codec(
{
RoundType.UPWARD: "upward",
RoundType.TOEVEN: "even",
RoundType.TRUNC: "truncate",
RoundType.TONEAREST: "nearest"
},
"RoundType"
)
[docs]
COLORSPACESTANDARD_CODEC: Codec[EncState, DecState, ColorSpaceStandard] = enum_codec(
{
ColorSpaceStandard.BT601: "bt601",
ColorSpaceStandard.BT709: "bt709",
ColorSpaceStandard.BT2020: "bt2020"
},
"ColorSpaceStandard"
)
[docs]
COLORCONVERSIONDIRECTION_CODEC: Codec[EncState, DecState, ColorConversion] = enum_codec(
{
ColorConversion.YUV2RGB: "yuv2rgb",
ColorConversion.RGB2YUV: "rgb2yuv",
ColorConversion.BGR2RGB: "bgr2rgb",
ColorConversion.RGB2BGR: "rgb2bgr"
},
"ColorConversionDirection"
)
[docs]
CHROMASUBSAMPLING_CODEC: Codec[EncState, DecState, ChromaSampling] = enum_codec(
{
ChromaSampling.NV12: "nv12",
ChromaSampling.YUV420: "yuv420",
ChromaSampling.YUV422: "yuv422"
},
"ChromaSubSampling"
)
[docs]
RESIZEMETHOD_CODEC: Codec[EncState, DecState, ResizeMethod] = enum_codec(
{
ResizeMethod.LINEAR: "linear",
ResizeMethod.NEAREST: "nearest",
ResizeMethod.AREA: "area",
ResizeMethod.CUBIC: "cubic"
},
"ResizeMethod"
)
[docs]
RESIZEDEPOSITLOCATION_CODEC: Codec[EncState, DecState, ResizeDepositLocation] = enum_codec(
{
ResizeDepositLocation.TOPLEFT: "topleft",
ResizeDepositLocation.CENTER: "center",
ResizeDepositLocation.BOTTOMRIGHT: "bottomright"
},
"ResizeDepositLocation"
)
# Some fields can have type either Tuple[A] or List[A]. Both types are
# encoded as a YAML sequence, so we need an extra tag to distinguish them.
# Codec[EncStateV, DecStateV, Union[Tuple[int, ...], List[int]]]
_tuple_or_list_scalar_codec = intrusive_tagged_union_codec(
'sequence',
{
'tuple': singleton_dict_codec('value', homogeneous_tuple_codec(scalar_codec), "Union of tuple and list"),
'list': singleton_dict_codec('value', _list_scalar_codec, "Union of tuple and list")
},
lambda x: 'tuple' if isinstance(x, tuple) else 'list'
)
[docs]
class NumpyEncState:
"""
Encoder state for encoding numpy arrays. Array contents are not encoded in yaml. Instead, arrays
are saved here, and a string that identifies the array is encoded in yaml.
"""
_numpy_dict: Dict[str, np.ndarray] # Collection of saved arrays, indexed by a generated name
_numpy_index: int # Index used to generate unique names
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._numpy_dict = dict()
self._numpy_index = 0
[docs]
def get_arrays(self) -> Dict[str, np.ndarray]:
"""
Get a dictionary of numpy arrays that have been saved, suitable for numpy's savez function.
The dictionary will be needed for decoding.
"""
return self._numpy_dict
[docs]
def save_array(self, array: np.ndarray) -> str:
"""
Save an array.
:return: Newly generated string that identifies the saved array.
"""
arr_str = "np_arr_{}".format(self._numpy_index)
self._numpy_dict[arr_str] = array
self._numpy_index += 1
return arr_str
[docs]
class NumpyDecState:
"""
Decoder state for decoding numpy arrays. Array contents are not encoded in yaml. Instead, arrays
are saved here, and array contents are retrieved during decoding.
"""
_numpy_dict: Dict[str, np.ndarray] # Collection of saved arrays
def __init__(self, *, numpy_dict: Dict[str, np.ndarray], **kwargs):
super().__init__(**kwargs)
self._numpy_dict = numpy_dict
[docs]
def load_array(self, key: str) -> np.ndarray:
"""
Load an array.
:param key: String that identifies the array.
:return: Array identified by the string.
"""
try:
return self._numpy_dict[key]
except KeyError as e:
raise DecodeError("Unable to find saved array: {}".format(key))
class _NDArrayCodec:
def encode(self, state: NumpyEncState, obj: np.ndarray) -> Any:
if not isinstance(obj, np.ndarray):
raise EncodeError("Unexpected data when encoding array")
return state.save_array(obj)
def decode(self, state: NumpyDecState, doc: Any) -> np.ndarray:
if not isinstance(doc, str):
raise DecodeError("Unexpected data when decoding array")
data = state.load_array(doc)
# Currently bf16 numpy array is saved with type of void16 instead of bf16 because bf16 is
# not a numpy native data type. To read back the array with the correct dtype, use the
# temporary solution to view the memory region as bf16.
if data.dtype.name == "void16":
data = data.view("bfloat16")
# 4 bits quantized weights (int4) are saved with type of void8.
if data.dtype.name == "void8":
data = data.view("int4")
return data
[docs]
ndarray_codec: Codec[NumpyEncState, NumpyDecState, np.ndarray] = _NDArrayCodec()
[docs]
def get_list_tag_name(x: _A) -> str:
"""
Get tag name of a list
:param x: list of objects
:return: tagName of the list
"""
assert isinstance(x, list)
if len(x) > 0 and isinstance(x[0], np.float32):
return 'listNPFloat32'
else:
return 'listScalar'
[docs]
def get_tag_name(x: _A) -> str:
"""
Get tag name of an object
:param x: Object
:return: tagName
"""
if isinstance(x, tuple):
return 'tuple'
if isinstance(x, SCALAR_TYPES):
return 'scalar'
if isinstance(x, list):
return 'list'
if isinstance(x, np.ndarray):
return 'ndarray'
if isinstance(x, np.float32):
return 'npFloat32'
if isinstance(x, np.int64):
return 'npInt64'
if isinstance(x, np.int32):
return 'npInt32'
# Serialization state during encoding
[docs]
class SerializerEncState(NumpyEncState):
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Serialization state during decoding
[docs]
class SerializerDecState(NumpyDecState):
def __init__(self, **kwargs):
super().__init__(**kwargs)
#####################################################
# Beginning of codecs for AwesomeNet data structures
#####################################################
[docs]
SCALARTYPE_CODEC = enum_codec(
{
ScalarType.int8: 'int8',
ScalarType.uint8: 'uint8',
ScalarType.int16: 'int16',
ScalarType.uint16: 'uint16',
ScalarType.int32: 'int32',
ScalarType.uint32: 'uint32',
ScalarType.int64: 'int64',
ScalarType.uint64: 'uint64',
ScalarType.float32: 'float32',
ScalarType.float16: 'float16',
ScalarType.float64: 'float64',
ScalarType.bool: 'bool',
ScalarType.bfloat16: 'bfloat16'
},
"ScalarType"
)
[docs]
TENSORTYPE_CODEC = dataclass_style_codec(
TensorType,
(
FieldEncoding('scalar', 'scalar', SCALARTYPE_CODEC),
FieldEncoding('shape', 'shape', homogeneous_tuple_codec(scalar_codec))
)
)
_list_tensor_type_codec = list_codec(TENSORTYPE_CODEC)
[docs]
STATUS_CODEC = enum_codec(
{
Status.RELAY: 'relay', Status.CALIBRATED: 'calibrated', Status.SIMA_QUANTIZED: 'sima_quantized'
},
"Status"
)
[docs]
TARGET_CODEC = enum_codec(
{
Platform.GEN1: 'gen1',
Platform.GEN2: 'gen2'
},
"Target"
)
[docs]
METRIC_CODEC = enum_codec(
{
Metric.l1: "l1",
Metric.l2: "l2",
Metric.mae: "mae",
Metric.mse: "mse",
Metric.psnr: "psnr",
Metric.mean: "mean",
Metric.std: "std"
},
"Metric"
)
[docs]
LAYERSTATS_CODEC = dataclass_style_codec(
LayerStats,
(
FieldEncoding('metric', 'metric', METRIC_CODEC),
FieldEncoding('error_value', 'error_value', scalar_codec)
)
)
[docs]
QUANTIZATION_PRECISION_CODEC = enum_codec(
{
QuantizationPrecision.INT_8: 'int8',
QuantizationPrecision.INT_16: 'int16',
QuantizationPrecision.BFLOAT_16: 'bfloat16',
QuantizationPrecision.BFLOAT_16_INT8_WEIGHTS: 'bfloat16_int8_weight',
QuantizationPrecision.BFLOAT_16_INT4_WEIGHTS: 'bfloat16_int4_weight'
},
"QuantizationPrecision"
)
[docs]
REQUANT_MODE_CODEC = enum_codec(
{
RequantizationMode.sima: 'sima',
RequantizationMode.tflite: 'tflite'
},
"RequantizationMode"
)
[docs]
BIAS_CORRECTION_TYPE_CODEC = enum_codec(
{
BiasCorrectionType.NONE: 'none',
BiasCorrectionType.REGULAR: 'regular',
BiasCorrectionType.ITERATIVE: 'iterative'
},
"BiasCorrectionType"
)
[docs]
empty_val_codec = dataclass_style_codec(
EmptyValue,
(FieldEncoding('empty', 'empty', scalar_codec),)
)
[docs]
def opt_codec(value_codec: Codec[EncState, DecState, _A]):
codec = dataclass_style_codec(
Opt,
(FieldEncoding('value', 'value', intrusive_tagged_union_codec(
'Option', {
'empty': singleton_dict_codec('no_value_set', empty_val_codec, 'Union of empty and any value.'),
'value': singleton_dict_codec('has_value', value_codec, 'Union of empty and any value.')
},
lambda x: 'empty' if isinstance(x, EmptyValue) else 'value'), ), ))
return codec
[docs]
QUANTIZATION_CONFIGS_CODEC = dataclass_style_codec(
QuantizationConfigs,
(
FieldEncoding('asymmetry', 'asymmetric', opt_codec(scalar_codec)),
FieldEncoding('per_channel', 'per_channel', opt_codec(scalar_codec)),
FieldEncoding('leaky_relu_uses_udf', 'leaky_relu_uses_udf', opt_codec(scalar_codec)),
FieldEncoding('quantization_precision', 'quantization_precision', opt_codec(QUANTIZATION_PRECISION_CODEC)),
FieldEncoding('intermediate_int32', 'intermediate_int32', opt_codec(scalar_codec)),
FieldEncoding('biascorr_type', 'bias_correction', opt_codec(BIAS_CORRECTION_TYPE_CODEC)),
FieldEncoding('requantization_mode', 'requantization_mode', opt_codec(REQUANT_MODE_CODEC)),
FieldEncoding('output_int32', 'output_int32', opt_codec(scalar_codec)),
)
)
[docs]
QUANTIZATION_CODEC = dataclass_style_codec(
Quantization,
(
FieldEncoding('scale', 'scale', scalar_codec),
FieldEncoding('zero_point', 'zero_point', scalar_codec),
FieldEncoding('bits', 'bits', scalar_codec, 8),
FieldEncoding('min_val', 'min_val', scalar_codec, -128.0),
FieldEncoding('max_val', 'max_val', scalar_codec, 127.0),
)
)
[docs]
REQUANT_METHOD_CODEC = enum_codec(
{
RequantMethod.fractional_zero: 'fractional_zero',
RequantMethod.arith_folded: 'arith_folded',
RequantMethod.scaled_fz: 'scaled_fractional_zero'
},
"RequantMethod"
)
[docs]
QUANT_RESULT_TENSOR_TYPE_CODEC = dataclass_style_codec(
QuantResultTensorType,
(
FieldEncoding('type', 'type', TENSORTYPE_CODEC),
FieldEncoding('quant', 'quant', optional_codec(QUANTIZATION_CODEC)),
FieldEncoding('requant_method', 'requant_method', optional_codec(REQUANT_METHOD_CODEC))
)
)
# Numpy dtype is converted to ScalarType for serialization
[docs]
DTYPE_CODEC: Codec[SerializerEncState, SerializerDecState, np.dtype] = \
map_codec(lambda s: np.dtype(s.numpy_type()), ScalarType.from_numpy, SCALARTYPE_CODEC)
def _unwrap_arith_folded_requantization(r: requantization.ArithFoldedRequantization) -> Dict[str, Any]:
return {'shift': r.shift, 'rounding': r.rounding, 'out_dtype': r.out_dtype}
def _wrap_arith_folded_requantization(d: Dict[str, Any]) -> requantization.ArithFoldedRequantization:
return requantization.narrowing_requantization(d['shift'], d['rounding'], d['out_dtype'])
_ARITH_FOLDED_REQUANTIZATION_CODEC = map_codec(
_wrap_arith_folded_requantization,
_unwrap_arith_folded_requantization,
named_dict_codec({'shift': _scalar_or_ndarray_codec, 'rounding': ROUNDTYPE_CODEC, 'out_dtype': DTYPE_CODEC},
description="ArithFoldedRequantization",
defaults={'out_dtype': np.dtype('int8')})
)
def _unwrap_fractional_zero_requantization(r: requantization.FractionalZeroRequantization) -> Dict[str, Any]:
return {'sc_correction': r.sc_correction, 'zp_correction': r.zp_correction,
'shift': r.narrowing.shift, 'rounding': r.narrowing.rounding, 'out_dtype': r.narrowing.out_dtype}
def _wrap_fractional_zero_requantization(d: Dict[str, Any]) -> requantization.FractionalZeroRequantization:
narrowing = requantization.Narrowing(d['shift'], d['rounding'], d['out_dtype'])
return requantization.FractionalZeroRequantization(d['sc_correction'], d['zp_correction'], narrowing)
_FRACTIONAL_ZERO_REQUANTIZATION_CODEC = map_codec(
_wrap_fractional_zero_requantization,
_unwrap_fractional_zero_requantization,
named_dict_codec({'sc_correction': scalar_codec, 'zp_correction': scalar_codec,
'shift': _scalar_or_ndarray_codec, 'rounding': ROUNDTYPE_CODEC, 'out_dtype': DTYPE_CODEC},
description="FractionalZeroRequantization",
defaults={'out_dtype': np.dtype('int8')})
)
def _unwrap_renormalization(r: requantization.Renormalization) -> Dict[str, Any]:
return {'sc_correction': r.sc_correction,
'shift': r.narrowing.shift, 'rounding': r.narrowing.rounding, 'out_dtype': r.narrowing.out_dtype}
def _wrap_renormalization(d: Dict[str, Any]) -> requantization.Renormalization:
narrowing = requantization.Narrowing(d['shift'], d['rounding'], d['out_dtype'])
return requantization.Renormalization(d['sc_correction'], narrowing)
_RENORMALIZATION_CODEC = map_codec(
_wrap_renormalization,
_unwrap_renormalization,
named_dict_codec({'sc_correction': _scalar_or_ndarray_codec,
'shift': _scalar_or_ndarray_codec, 'rounding': ROUNDTYPE_CODEC, 'out_dtype': DTYPE_CODEC},
description="Renormalization",
defaults={'out_dtype': np.dtype('bfloat16')})
)
[docs]
REQUANTIZATION_CODEC: Codec[SerializerEncState, SerializerDecState, requantization.BaseRequantization[np.ndarray]] = adt_style_codec(
'requant_type',
(
ConstructorEncoding(requantization.TFLiteRequantization, 'tflite', (
FieldEncoding('sc_correction', 'sc_correction', _scalar_or_ndarray_codec),
FieldEncoding('shift', 'shift', _scalar_or_ndarray_codec),
FieldEncoding('rounding', 'rounding', ROUNDTYPE_CODEC),
FieldEncoding('zp_correction', 'zp_correction', scalar_codec),
FieldEncoding('out_dtype', 'out_dtype', DTYPE_CODEC, np.dtype('int8'))
)),
ConstructorEncoding(requantization.ArithFoldedRequantization, 'arith_folded',
_ARITH_FOLDED_REQUANTIZATION_CODEC),
ConstructorEncoding(requantization.FractionalZeroRequantization, 'fractional_zero',
_FRACTIONAL_ZERO_REQUANTIZATION_CODEC),
ConstructorEncoding(requantization.Renormalization, 'renormalize',
_RENORMALIZATION_CODEC),
)
)
[docs]
BACKEND_CODEC = enum_codec({value: name for (name, value) in Backend.__members__.items()}, "Backend")
[docs]
def data_value_codec(c: Codec[SerializerEncState, SerializerDecState, _A]) \
-> Codec[SerializerEncState, SerializerDecState, DataValue[_A]]:
ret = lazy_codec(lambda: adt_style_codec(
'type',
(
ConstructorEncoding(TensorValue, 'tensor', (FieldEncoding('value', 'value', c),)),
ConstructorEncoding(TupleValue, 'tuple', (FieldEncoding('elements', 'elements', list_codec(ret)),))
)
))
return ret
[docs]
def node_associated_value_codec(c: Codec[SerializerEncState, SerializerDecState, _A]) \
-> Codec[SerializerEncState, SerializerDecState, NodeAssociatedValue[_A]]:
return dataclass_style_codec(NodeAssociatedValue, (
FieldEncoding('inputs', 'inputs', dict_codec(scalar_codec, data_value_codec(c))),
FieldEncoding('output', 'output', data_value_codec(c))
))
def _wrap_simairmetadata(doc: Dict[str, Any]) -> SiMaIRMetadata:
return SiMaIRMetadata(doc['relay_ir_name'])
def _unwrap_simairmetadata(obj: SiMaIRMetadata) -> Dict[str, Any]:
return {'relay_ir_name': obj['relay_ir_name']}
# We treat AwesomeOperation types (the classes) as enum values.
# Item order follows the order in sima_ir.py.
_AWESOMEOPERATION_TYPE_CODEC = enum_codec(
{
operations.PlaceholderOp: 'placeholder',
operations.ConstantOp: 'constant',
operations.MaxPool2DOp: 'max_pool2d',
operations.AvgPool2DOp: 'avg_pool2d',
operations.AdaptiveAvgPool2DOp: 'adaptive_avg_pool2d',
operations.MaxPool3DOp: 'max_pool3d',
operations.AvgPool3DOp: 'avg_pool3d',
operations.MultiplyOp: 'multiply',
operations.MeanOp: 'mean',
operations.SqueezeOp: 'squeeze',
operations.ArgMaxOp: 'argmax',
operations.SoftmaxOp: 'softmax',
operations.TupleOp: 'tuple',
operations.PadOp: 'pad',
operations.LRNOp: 'lrn',
operations.LayerNormOp: 'layer_norm',
operations.RMSNormOp: 'rms_norm',
operations.ConcatenateOp: 'concat',
operations.TransposeOp: 'transpose',
operations.DepthToSpaceOp: 'depth_to_space',
operations.ReshapeOp: 'reshape',
operations.ExpandDimsOp: 'expand_dims',
operations.UpsamplingOp: 'upsample',
operations.LayoutTransformOp: 'layout_transform',
operations.CastOp: 'cast',
operations.ExtmOp: 'extm',
operations.SumOp: 'sum',
operations.ProdOp: 'prod',
operations.TupleGetItemOp: 'tuple_get_item',
operations.SubtractOp: 'subtract',
operations.PowerOp: 'power',
operations.MaximumOp: 'maximum',
operations.MinimumOp: 'minimum',
operations.FullOp: 'full',
operations.TileOp: 'tile',
operations.SplitOp: 'split',
operations.TakeOp: 'take',
operations.StridedSliceOp: 'strided_slice',
operations.ImageResize2DOp: 'image_resize2d',
operations.SqrtOp: 'sqrt',
operations.RsqrtOp: 'rsqrt',
operations.TanhOp: 'tanh',
operations.SigmoidOp: 'sigmoid',
operations.LogOp: 'log',
operations.Log2Op: 'log2',
operations.Log10Op: 'log10',
operations.DivideOp: 'divide',
operations.LeakyReluCompositeOp: 'leaky_relu_udf',
operations.EluOp: 'elu',
operations.ErfOp: 'erf',
operations.PReluOp: 'prelu',
operations.ReluOp: 'relu',
operations.ExpOp: 'exp',
operations.ConvAddActivationOp: 'conv_add_activation',
operations.AddActivationOp: 'add_activation',
operations.TupleConcatenateOp: 'tuple_concat',
operations.SwishOp: 'swish',
operations.HardSwishOp: 'hard_swish',
operations.HardSigmoidOp: 'hard_sigmoid',
operations.ConstantMultiplyAddOp: 'constant_multiply_add',
operations.ExternalOp: 'external',
operations.QNNQuantizeOp: 'qnn_quantize',
operations.RequantizeOp: 'qnn_requantize',
operations.QNNDequantizeOp: 'qnn_dequantize',
operations.QNNMulOp: 'qnn_mul',
operations.SigmoidTransformOp: 'sigmoid_transform',
operations.NmsMaxpoolTransformOp: 'nms_maxpool',
operations.QuantizationTransformOp: 'quantize',
operations.DequantizationTransformOp: 'dequantize',
operations.ChromaUpsampleTransformOp: 'chromaupsample',
operations.YuvRgbConversionTransformOp: 'yuvrgb',
operations.BgrRgbConversionTransformOp: 'bgrrgb',
operations.NormalizationTransformOp: 'normalization',
operations.ResizeTransformOp: 'resize',
operations.SoftplusOp: 'softplus',
operations.BatchMatmulOp: 'batch_matmul',
operations.UnaryBatchMatmulOp: 'unary_batch_matmul',
operations.SliceConcatOp: 'slice_concat',
operations.BroadcastToOp: 'broadcast_to',
operations.GeluOp: 'gelu',
operations.InstanceNormOp: 'instance_norm',
operations.GridSampleOp: 'grid_sample',
operations.VarianceOp: 'variance',
operations.ClipOp: 'clip',
},
"Operation"
)
# AwesomeOperation instances don't contain any data besides the type, so we encode instances by
# just encoding the type.
[docs]
AWESOMEOPERATION_CODEC = map_codec(lambda cls: cls(), lambda inst: type(inst), _AWESOMEOPERATION_TYPE_CODEC)
##################################################
# Attributes and attribute fields that are reused
##################################################
_constant_attrs_codec = dataclass_style_codec(attributes.ConstantAttrs,
(FieldEncoding('data', 'data', ndarray_codec),))
_conv_attrs_codec = dataclass_style_codec(attributes.ConvAttrs, (
FieldEncoding('stride', 'stride', homogeneous_tuple_codec(scalar_codec)),
FieldEncoding('dilation', 'dilation', homogeneous_tuple_codec(scalar_codec)),
FieldEncoding('padding', 'padding', homogeneous_tuple_codec(homogeneous_tuple_codec(scalar_codec))),
FieldEncoding('output_padding', 'output_padding', homogeneous_tuple_codec(homogeneous_tuple_codec(scalar_codec))),
FieldEncoding('is_transposed', 'is_transposed', scalar_codec, False),
FieldEncoding('weight_shape', 'weight_shape', homogeneous_tuple_codec(scalar_codec)),
FieldEncoding('input_spatial_shape', 'input_spatial_shape', homogeneous_tuple_codec(scalar_codec)),
FieldEncoding('batch_size', 'batch_size', scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC)
))
# The common fields of pooling attributes.
_POOL_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('ceil_mode', 'ceil_mode', scalar_codec),
FieldEncoding('out_layout', 'out_layout', scalar_codec),
FieldEncoding('layout', 'layout', scalar_codec),
FieldEncoding('padding', 'padding', homogeneous_tuple_codec(homogeneous_tuple_codec(scalar_codec))),
FieldEncoding('pool_size', 'pool_size', _list_scalar_codec),
FieldEncoding('strides', 'strides', _list_scalar_codec),
FieldEncoding('dilation', 'dilation', _list_scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC)
)
# The fields of MaxPool2DAttrs and MaxPool3DAttrs.
_MAXPOOL_ATTRS_FIELDS: Sequence[FieldEncoding] = _POOL_ATTRS_FIELDS
# The fields of AvgPool2DAttrs and AvgPool3DAttrs.
_AVGPOOL_ATTRS_FIELDS: Sequence[FieldEncoding] = (
*_POOL_ATTRS_FIELDS,
FieldEncoding('count_include_pad', 'count_include_pad', scalar_codec)
)
# The fields of MaximumAttrs and MinimumAttrs.
_MAXIMUM_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)
_maximum_attrs_codec = dataclass_style_codec(attributes.MaximumAttrs, _MAXIMUM_ATTRS_FIELDS)
_minimum_attrs_codec = dataclass_style_codec(attributes.MinimumAttrs, _MAXIMUM_ATTRS_FIELDS)
# The common fields of reduction attributes, including mean and argmax.
_REDUCE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('axis', 'axis', _list_scalar_codec),
FieldEncoding('exclude', 'exclude', scalar_codec, False),
FieldEncoding('keepdims', 'keepdims', scalar_codec, False),
FieldEncoding('shape', 'shape', _tuple_scalar_codec)
)
# Fields of class TupleAttrs
_TUPLE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_types', 'input_types', _list_tensor_type_codec),
)
_tuple_attrs_codec = dataclass_style_codec(attributes.TupleAttrs, _TUPLE_ATTRS_FIELDS)
# Fields of class ConcatenateAttrs
_CONCATENATE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('input_types', 'input_types', _list_tensor_type_codec)
)
_concatenate_attrs_codec = dataclass_style_codec(attributes.ConcatenateAttrs, _CONCATENATE_ATTRS_FIELDS)
# Fields of TupleConcatenateAttrs
_TUPLE_CONCAT_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('tuple_attrs', 'tuple', _tuple_attrs_codec),
FieldEncoding('concat_attrs', 'concat', _concatenate_attrs_codec),
)
_tuple_concat_attrs_codec = dataclass_style_codec(attributes.TupleConcatenateAttrs, _TUPLE_CONCAT_ATTRS_FIELDS)
# Fields of class TransposeAttrs
_TRANSPOSE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('axes', 'axes', _list_scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC),
)
_transpose_attrs_codec = dataclass_style_codec(attributes.TransposeAttrs, _TRANSPOSE_ATTRS_FIELDS)
# Fields of class DepthToSpaceAttrs
_DEPTHTOSPACE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('block_size', 'block_size', scalar_codec),
FieldEncoding('mode', 'mode', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC),
)
_depth_to_space_attrs_codec = dataclass_style_codec(attributes.DepthToSpaceAttrs, _DEPTHTOSPACE_ATTRS_FIELDS)
# Fields of UDF attributes class
_UDF_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
)
_udf_attrs_codec = dataclass_style_codec(attributes.UDFAttrs, _UDF_ATTRS_FIELDS)
# Fields of class UDFAQuantattrs
_UDF_QUANT_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('attrs', 'attrs', _udf_attrs_codec),
FieldEncoding('input_signed', 'input_signed', scalar_codec, False),
FieldEncoding('output_signed', 'output_signed', scalar_codec, False),
FieldEncoding('lookup_table', 'lookup_table', optional_codec(ndarray_codec)),
FieldEncoding('input_int16', 'input_int16', scalar_codec),
FieldEncoding('requant', 'requant', optional_codec(REQUANTIZATION_CODEC))
)
_udf_quant_attrs_codec = dataclass_style_codec(attributes.UDFQuantAttrs, _UDF_QUANT_ATTRS_FIELDS)
# Fields of class MultiplyQuantattrs
_MULTIPLY_QUANT_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('lhs_input_shape', 'lhs_input_shape', _tuple_scalar_codec),
FieldEncoding('rhs_input_shape', 'rhs_input_shape', _tuple_scalar_codec),
FieldEncoding('input_int16', 'input_int16', scalar_codec),
FieldEncoding('intrinsic_shift', 'intrinsic_shift', scalar_codec),
FieldEncoding('requant', 'requant', REQUANTIZATION_CODEC),
FieldEncoding('lhs_zero_point', 'lhs_zero_point', scalar_codec, 0),
FieldEncoding('rhs_zero_point', 'rhs_zero_point', scalar_codec, 0),
FieldEncoding('layer_bits', 'layer_bits', scalar_codec, 8),
)
_multply_quant_attrs_codec = dataclass_style_codec(attributes.MultiplyQuantAttrs, _MULTIPLY_QUANT_ATTRS_FIELDS)
_LAYER_NORM_QUANT_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('axis', 'axis', _scalar_or_tuple_scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('zp_rsqrt', 'zp_rsqrt', scalar_codec),
FieldEncoding('lookup_table_rsqrt', 'lookup_table_rsqrt', optional_codec(ndarray_codec)),
FieldEncoding('requant_mean', 'requant_mean', REQUANTIZATION_CODEC),
FieldEncoding('requant_lut_input', 'requant_lut_input', REQUANTIZATION_CODEC),
FieldEncoding('requant_output', 'requant_output', REQUANTIZATION_CODEC),
)
_layer_norm_quant_attrs_codec = dataclass_style_codec(attributes.LayerNormQuantAttrs, _LAYER_NORM_QUANT_ATTRS_FIELDS)
_EXTM_ATTRS_FILDS = (
*_REDUCE_ATTRS_FIELDS,
FieldEncoding('max', 'max', scalar_codec)
)
_extm_attrs_codec = dataclass_style_codec(attributes.ExtmAttrs, _EXTM_ATTRS_FILDS)
_ARGMAX_ATTRS_FIELDS: Sequence[FieldEncoding] = (
*_REDUCE_ATTRS_FIELDS,
FieldEncoding('select_last_index', 'select_last_index', scalar_codec, False),
FieldEncoding('result_scalar_type', 'result_type', SCALARTYPE_CODEC),
FieldEncoding('input_scalar_type', 'input_type', SCALARTYPE_CODEC)
)
_argmax_attrs_codec = dataclass_style_codec(attributes.ArgMaxAttrs, _ARGMAX_ATTRS_FIELDS)
# Fields of class ClipAttrs
_CLIP_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('a_min', 'min', scalar_codec),
FieldEncoding('a_max', 'max', scalar_codec),
FieldEncoding('shape', 'shape', _tuple_scalar_codec),
FieldEncoding("scalar_type", "scalar_type", SCALARTYPE_CODEC),
)
_clip_attrs_codec = dataclass_style_codec(attributes.ClipAttrs, _CLIP_ATTRS_FIELDS)
_relu_attrs_codec = dataclass_style_codec(attributes.ReluAttrs, (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
))
_relu_quant_attrs_codec = dataclass_style_codec(attributes.ReluQuantAttrs, (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('zero_point', 'zero_point', scalar_codec),
))
_ADD_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('lhs_input_shape', 'lhs_input_shape', _tuple_scalar_codec),
FieldEncoding('rhs_input_shape', 'rhs_input_shape', _tuple_scalar_codec)
)
_add_attrs_codec = dataclass_style_codec(attributes.AddAttrs, _ADD_ATTRS_FIELDS)
# Fields of class BiasAddAttrs
_BIAS_ADD_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('axis', 'axis', scalar_codec)
)
_bias_add_attrs_codec = dataclass_style_codec(attributes.BiasAddAttrs, _BIAS_ADD_ATTRS_FIELDS)
# In a fused convolution or matrix multiply, the add operator can be either of these
_add_or_bias_add_attrs_codec = type_tagged_codec(
'add_type',
(
('per_channel', attributes.BiasAddAttrs, _bias_add_attrs_codec),
('per_element', attributes.AddAttrs, _add_attrs_codec)
))
_UPSAMPLING_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('scale_h', 'scale_h', scalar_codec),
FieldEncoding('scale_w', 'scale_w', scalar_codec),
FieldEncoding('layout', 'layout', scalar_codec),
FieldEncoding('method', 'method', scalar_codec),
FieldEncoding('align_corners', 'align_corners', scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC)
)
_upsampling_attrs_codec = dataclass_style_codec(attributes.UpsamplingAttrs, _UPSAMPLING_ATTRS_FIELDS)
def _conv_add_activation_fields() \
-> List[FieldEncoding]:
"""
Create the list of fields of ConvAddActivationAttrs.
"""
return list(filter(lambda x: x is not None, (
FieldEncoding('weights_attrs', 'weights', _constant_attrs_codec),
FieldEncoding('conv_attrs', 'conv', _conv_attrs_codec),
FieldEncoding('bias_attrs', 'bias', optional_codec(_constant_attrs_codec)),
FieldEncoding('add_attrs', 'add', optional_codec(_add_or_bias_add_attrs_codec)),
FieldEncoding('activ_attrs', 'activ', optional_codec(_activ_attrs_codec))
)))
# Activation attributes can be ClipAttrs or ReluAttrs
_activ_attrs_codec = type_tagged_codec(
'activation',
(
('clip', attributes.ClipAttrs, _clip_attrs_codec),
('relu', attributes.ReluAttrs, _relu_attrs_codec)
))
# Activation attributes can be ClipAttrs or ReluAttrs
_activ_quant_attrs_codec = type_tagged_codec(
'activation',
(
('fp_clip', attributes.ClipAttrs, _clip_attrs_codec),
('fp_relu', attributes.ReluAttrs, _relu_attrs_codec),
('clip', attributes.ClipQuantAttrs, _clip_attrs_codec),
('relu', attributes.ReluQuantAttrs, _relu_quant_attrs_codec)
))
def _add_activation_fields() -> List[FieldEncoding]:
"""
Create the list of fields of AddActivationAttrs.
:return: Fields of the class
"""
return list(filter(lambda x: x is not None, (
FieldEncoding('add_attrs', 'add', _add_attrs_codec),
FieldEncoding('activ_attrs', 'activation', optional_codec(_activ_attrs_codec))
)))
# The fields of LeakyReluQuantAttrs
_LEAKY_RELU_QUANT_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('alpha', 'alpha', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('right_shift', 'right_shift', scalar_codec),
FieldEncoding('zero_point', 'zero_point', scalar_codec),
FieldEncoding('bits', 'bits', scalar_codec),
FieldEncoding('rounding_type', 'rounding', ROUNDTYPE_CODEC)
)
_leaky_relu_quant_attrs_codec = dataclass_style_codec(attributes.LeakyReluQuantAttrs, _LEAKY_RELU_QUANT_ATTRS_FIELDS)
# The fields of ConcatQuantAttrs
_CONCAT_QUANT_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('attrs', 'attrs', _concatenate_attrs_codec),
FieldEncoding('requants', 'requants', list_codec(REQUANTIZATION_CODEC)),
FieldEncoding('layer_bits', 'layer_bits', _list_scalar_codec, [8]),
FieldEncoding('input_scales', 'input_scales', list_codec(_scalar_or_list_scalar_codec)),
FieldEncoding('node_scales', 'node_scales', _list_scalar_codec),
FieldEncoding('node_zps', 'node_zps', _list_scalar_codec),
)
_concat_quant_attrs_codec = dataclass_style_codec(attributes.ConcatQuantAttrs, _CONCAT_QUANT_ATTRS_FIELDS)
# The fields of StridedSliceAttrs
_STRIDED_SLICE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('begin', 'begin', _list_scalar_codec),
FieldEncoding('end', 'end', _list_scalar_codec),
FieldEncoding('strides', 'strides', _list_scalar_codec),
FieldEncoding('axes', 'axes', optional_codec(_tuple_or_list_scalar_codec)),
FieldEncoding('slice_mode', 'slice_mode', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC)
)
_strided_slice_attrs_codec = dataclass_style_codec(attributes.StridedSliceAttrs, _STRIDED_SLICE_ATTRS_FIELDS)
_BROADCAST_TO_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
FieldEncoding('output_shape', 'output_shape', _tuple_scalar_codec),
)
_broadcast_to_attrs_codec = dataclass_style_codec(attributes.BroadcastToAttrs, _BROADCAST_TO_ATTRS_FIELDS)
_MULTIPLY_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('lhs_input_shape', 'lhs_input_shape', _tuple_scalar_codec),
FieldEncoding('rhs_input_shape', 'rhs_input_shape', _tuple_scalar_codec),
)
_multiply_attrs_codec = dataclass_style_codec(attributes.MultiplyAttrs, _MULTIPLY_ATTRS_FIELDS)
_LAYER_NORM_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('axis', 'axis', _scalar_or_tuple_scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('epsilon', 'epsilon', scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC)
)
_layer_norm_attrs_codec = dataclass_style_codec(attributes.LayerNormAttrs, _LAYER_NORM_ATTRS_FIELDS)
_INSTANCE_NORM_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('axis', 'axis', _tuple_scalar_codec),
FieldEncoding('input_data_shape', 'input_data_shape', _tuple_scalar_codec),
FieldEncoding('mean_shape', 'mean_shape', _tuple_scalar_codec),
FieldEncoding('variance_shape', 'variance_shape', _tuple_scalar_codec),
FieldEncoding('epsilon', 'epsilon', scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC)
)
_instance_norm_attrs_codec = dataclass_style_codec(attributes.InstanceNormAttrs, _INSTANCE_NORM_ATTRS_FIELDS)
_GRID_SAMPLE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('grid_shape', 'grid_shape', _tuple_scalar_codec),
FieldEncoding('method', 'method', scalar_codec),
FieldEncoding('padding_mode', 'padding_mode', scalar_codec),
FieldEncoding('align_corners', 'align_corners', scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC)
)
_grid_sample_attrs_codec = dataclass_style_codec(attributes.GridSampleAttrs, _GRID_SAMPLE_ATTRS_FIELDS)
_VARIANCE_ATTRS_FIELDS: Sequence[FieldEncoding] = (
FieldEncoding('input_data_shape', 'input_data_shape', _tuple_scalar_codec),
FieldEncoding('mean_shape', 'mean_shape', _tuple_scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('axis', 'axis', _tuple_scalar_codec)
)
_variance_attrs_codec = dataclass_style_codec(attributes.VarianceAttrs, _VARIANCE_ATTRS_FIELDS)
def _wrap_external_attrs(doc: dict) -> attributes.ExternalAttrs:
"""Construct an ExternalAttrs after deserializing its data."""
return attributes.ExternalAttrs(doc['external_input_list'], doc['type'], doc['backend'],
doc['irmodule'], doc['operations'])
def _unwrap_external_attrs(obj: attributes.ExternalAttrs) -> dict:
"""Extract from an ExternalAttrs the data to serialize."""
return {'external_input_list': obj.external_input_list, 'type': obj.node_type,
'backend': obj.backend, 'irmodule': obj.irmod_str,
'operations': obj.operations}
# Codec for ExternalAttrs.
# Some of the class fields are TVM object references that we don't include in the serialized representation.
# These fields are generated from the other fields when needed.
_external_attrs_codec = map_codec(
_wrap_external_attrs,
_unwrap_external_attrs,
named_dict_codec(
{
'external_input_list': _list_scalar_codec,
'type': node_associated_value_codec(TENSORTYPE_CODEC),
'backend': scalar_codec,
'irmodule': scalar_codec,
'operations': _list_scalar_codec
},
"ExternalAttrs"
))
# Codec for AwesomeAttributes.
# Item order follows the order in sima_ir.py.
[docs]
AWESOMEATTRIBUTES_CODEC = adt_style_codec(
'op',
[
ConstructorEncoding(attributes.PlaceholderAttrs, 'placeholder', (
FieldEncoding('type', 'type', TENSORTYPE_CODEC),
)),
ConstructorEncoding(attributes.ConstantAttrs, 'constant', _constant_attrs_codec),
ConstructorEncoding(attributes.MaxPoolAttrs, 'max_pool', _MAXPOOL_ATTRS_FIELDS),
ConstructorEncoding(attributes.AvgPoolAttrs, 'avg_pool', _AVGPOOL_ATTRS_FIELDS),
ConstructorEncoding(attributes.AdaptiveAvgPool2DAttrs, 'adaptive_avg_pool2d', (
FieldEncoding('output_size', 'output_size', homogeneous_tuple_codec(scalar_codec)),
FieldEncoding('out_layout', 'out_layout', scalar_codec),
FieldEncoding('layout', 'layout', scalar_codec)
)),
ConstructorEncoding(attributes.ReluAttrs, 'relu', (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.AddAttrs, 'add', _ADD_ATTRS_FIELDS),
ConstructorEncoding(attributes.BiasAddAttrs, 'bias_add', _BIAS_ADD_ATTRS_FIELDS),
ConstructorEncoding(attributes.AwesomeAttributes, 'multiply', ()),
ConstructorEncoding(attributes.MeanAttrs, 'mean', _REDUCE_ATTRS_FIELDS),
ConstructorEncoding(attributes.SqueezeAttrs, 'squeeze', (
FieldEncoding('axis', 'axis', _list_scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC),
)),
ConstructorEncoding(attributes.ArgMaxAttrs, 'argmax', _ARGMAX_ATTRS_FIELDS),
ConstructorEncoding(attributes.SoftmaxAttrs, 'softmax', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
)),
ConstructorEncoding(attributes.TupleAttrs, 'tuple', _tuple_attrs_codec),
ConstructorEncoding(attributes.PadAttrs, 'pad', (
FieldEncoding('pad_mode', 'pad_mode', scalar_codec),
FieldEncoding('pad_width', 'pad_width', homogeneous_tuple_codec(homogeneous_tuple_codec(scalar_codec))),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec)
)),
ConstructorEncoding(attributes.LRNAttrs, 'lrn', (
FieldEncoding('alpha', 'alpha', scalar_codec),
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('beta', 'beta', scalar_codec),
FieldEncoding('bias', 'bias', scalar_codec),
FieldEncoding('size', 'size', scalar_codec),
FieldEncoding('shape', 'shape', _tuple_scalar_codec)
)),
ConstructorEncoding(attributes.ConcatenateAttrs, 'concat', _concatenate_attrs_codec),
ConstructorEncoding(attributes.TransposeAttrs, 'transpose', _TRANSPOSE_ATTRS_FIELDS),
ConstructorEncoding(attributes.DepthToSpaceAttrs, 'depth_to_space', _DEPTHTOSPACE_ATTRS_FIELDS),
ConstructorEncoding(attributes.ClipAttrs, 'clip', _CLIP_ATTRS_FIELDS),
ConstructorEncoding(attributes.ReshapeAttrs, 'reshape', (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('dtype', 'dtype', SCALARTYPE_CODEC),
FieldEncoding('newshape', 'output_shape', _list_scalar_codec),
)),
ConstructorEncoding(attributes.ExpandDimsAttrs, 'expand_dims', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('num_newaxis', 'num_newaxis', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC),
)),
ConstructorEncoding(attributes.BatchFlattenAttrs, 'batch_flatten', ()),
ConstructorEncoding(attributes.UpsamplingAttrs, 'upsample', _UPSAMPLING_ATTRS_FIELDS),
ConstructorEncoding(attributes.LayoutTransformAttrs, 'layout_transform', (
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
FieldEncoding('src_layout', 'src_layout', scalar_codec),
FieldEncoding('dst_layout', 'dst_layout', scalar_codec),
FieldEncoding('implicitly_removable', 'implicitly_removable', scalar_codec, False)
)),
ConstructorEncoding(attributes.TessellationTransformAttrs, 'tessellation_transform', (
FieldEncoding('slice_shape', 'slice_shape', _tuple_scalar_codec),
FieldEncoding('align_c16', 'align_c16', scalar_codec),
FieldEncoding('cblock', 'cblock', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.DetessellationTransformAttrs, 'detessellation_transform', (
FieldEncoding('slice_shape', 'slice_shape', _tuple_scalar_codec),
FieldEncoding('align_c16', 'align_c16', scalar_codec),
FieldEncoding('cblock', 'cblock', scalar_codec),
FieldEncoding('frame_type', 'frame_type', TENSORTYPE_CODEC),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.PackTransformAttrs, 'pack_transform', (
FieldEncoding('input_shapes', 'input_shapes', _list_tensor_type_codec),
FieldEncoding('result_scalar_type', 'result_type', SCALARTYPE_CODEC),
)),
ConstructorEncoding(attributes.UnpackTransformAttrs, 'unpack_transform', (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('tensor_types', 'tensor_types', _list_tensor_type_codec),
)),
ConstructorEncoding(attributes.NormalizationTransformAttrs, 'normalization_transform', (
FieldEncoding('channel_params', 'channel_params', _list_tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
)),
ConstructorEncoding(attributes.QuantizationTransformAttrs, 'quantization_transform', (
FieldEncoding('channel_params', 'channel_params', _list_tuple_scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('num_bits', 'num_bits', scalar_codec),
FieldEncoding('rounding', 'rounding', ROUNDTYPE_CODEC),
FieldEncoding('output_data_type', 'output_data_type', SCALARTYPE_CODEC)
)),
ConstructorEncoding(attributes.DequantizationTransformAttrs, 'dequantization_transform', (
FieldEncoding('channel_params', 'channel_params', _list_tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
FieldEncoding('output_type', 'output_type', SCALARTYPE_CODEC),
)),
ConstructorEncoding(attributes.ResizeTransformAttrs, 'resize_transform', (
FieldEncoding('target_height', 'target_height', scalar_codec),
FieldEncoding('target_width', 'target_width', scalar_codec),
FieldEncoding('keep_aspect', 'keep_aspect', scalar_codec),
FieldEncoding('deposit_location', 'deposit_location', RESIZEDEPOSITLOCATION_CODEC),
FieldEncoding('method', 'method', RESIZEMETHOD_CODEC),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
)),
ConstructorEncoding(attributes.ChromaUpsampleTransformAttrs, 'chroma_upsample_transform', (
FieldEncoding('frame_height', 'frame_height', scalar_codec),
FieldEncoding('frame_width', 'frame_width', scalar_codec),
FieldEncoding('yuv_sampling', 'yuv_sampling', CHROMASUBSAMPLING_CODEC),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
)),
ConstructorEncoding(attributes.YuvRgbConversionTransformAttrs, 'yuv_rgb_conversion_transform', (
FieldEncoding('conversion', 'conversion', COLORCONVERSIONDIRECTION_CODEC),
FieldEncoding('std', 'std', COLORSPACESTANDARD_CODEC),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.BgrRgbConversionTransformAttrs, 'bgr_rgb_conversion_transform', (
FieldEncoding('conversion', 'conversion', COLORCONVERSIONDIRECTION_CODEC),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.SigmoidTransformAttrs, 'sigmoid_transform', (
FieldEncoding('save_int16', 'save_int16', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.NmsMaxpoolTransformAttrs, 'nms_maxpool_transform', (
FieldEncoding('kernel', 'kernel', scalar_codec),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
)),
ConstructorEncoding(attributes.CastAttrs, 'cast', (
FieldEncoding('out_dtype', 'out_dtype', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC)
)),
ConstructorEncoding(attributes.ExtmAttrs, 'extm', (
*_REDUCE_ATTRS_FIELDS,
FieldEncoding('max', 'max', scalar_codec, False),
)),
ConstructorEncoding(attributes.SumAttrs, 'sum', (
*_REDUCE_ATTRS_FIELDS,
FieldEncoding('num_element', 'num_element', scalar_codec, 0),
)),
ConstructorEncoding(attributes.ProdAttrs, 'prod', _REDUCE_ATTRS_FIELDS),
ConstructorEncoding(attributes.TupleGetItemAttrs, 'tuple_get_item', (
FieldEncoding('input_types', 'input_types', _list_tensor_type_codec),
FieldEncoding('index', 'index', scalar_codec),
)),
ConstructorEncoding(attributes.SubtractAttrs, 'subtract', (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('lhs_input_shape', 'lhs_input_shape', _tuple_scalar_codec),
FieldEncoding('rhs_input_shape', 'rhs_input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.PowerAttrs, 'power', (
FieldEncoding('lhs_input_shape', 'lhs_input_shape', _tuple_scalar_codec),
FieldEncoding('rhs_input_shape', 'rhs_input_shape', _tuple_scalar_codec)
)),
ConstructorEncoding(attributes.MaximumAttrs, 'maximum', (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.MinimumAttrs, 'minimum', (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.FullAttrs, 'full', (
FieldEncoding('shape', 'shape', _list_scalar_codec),
FieldEncoding('dtype', 'dtype', scalar_codec)
)),
ConstructorEncoding(attributes.TileAttrs, 'tile', (
FieldEncoding('reps', 'reps', _list_scalar_codec),
)),
ConstructorEncoding(attributes.SplitAttrs, 'split', (
FieldEncoding('indices_or_sections', 'dividers', _scalar_or_tuple_scalar_codec),
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC),
)),
ConstructorEncoding(attributes.TakeAttrs, 'take', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('batch_dims', 'batch_dims', scalar_codec),
FieldEncoding('mode', 'mode', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('indices_shape', 'indices_shape', _tuple_scalar_codec),
FieldEncoding('input_type', 'input_type', SCALARTYPE_CODEC)
)),
ConstructorEncoding(attributes.StridedSliceAttrs, 'strided_slice', _strided_slice_attrs_codec),
ConstructorEncoding(attributes.BroadcastToAttrs, 'broadcast_to', _broadcast_to_attrs_codec),
ConstructorEncoding(attributes.ImageResize2DAttrs, 'resize2d', (
FieldEncoding('size', 'size', _list_scalar_codec),
FieldEncoding('roi', 'roi', homogeneous_tuple_codec(scalar_codec)),
FieldEncoding('layout', 'layout', scalar_codec),
FieldEncoding('method', 'method', scalar_codec),
FieldEncoding('coordinate_transformation_mode', 'coordinate_transformation_mode', scalar_codec),
FieldEncoding('rounding_method', 'rounding_method', scalar_codec),
FieldEncoding('cubic_alpha', 'cubic_alpha', scalar_codec),
FieldEncoding('cubic_exclude', 'cubic_exclude', scalar_codec),
FieldEncoding('extrapolation_value', 'extrapolation_value', scalar_codec),
FieldEncoding('out_dtype', 'out_dtype', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec)
)),
ConstructorEncoding(attributes.UDFAttrs, 'udf', _udf_attrs_codec),
ConstructorEncoding(attributes.DivideAttrs, 'divide', (
FieldEncoding('udf_attrs', 'udf_attrs', _udf_attrs_codec),
FieldEncoding('multiply_attrs', 'multiply_attrs', _multiply_attrs_codec),
)),
ConstructorEncoding(attributes.LeakyReluAttrs, 'leaky_relu', (
*_UDF_ATTRS_FIELDS,
FieldEncoding('alpha', 'alpha', scalar_codec),
)),
ConstructorEncoding(attributes.PReluAttrs, 'prelu', (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('alpha', 'alpha', ndarray_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.SwishAttrs, 'swish', _udf_attrs_codec),
ConstructorEncoding(attributes.ConvAddActivationAttrs, 'conv_add_activation', _conv_add_activation_fields()),
ConstructorEncoding(attributes.AddActivationAttrs, 'add_activation', _add_activation_fields()),
ConstructorEncoding(attributes.TupleConcatenateAttrs, 'tuple_concat', _tuple_concat_attrs_codec),
ConstructorEncoding(attributes.ConstantMultiplyAddAttrs, 'constant_multiply_add', (
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
FieldEncoding('lhs_input_shape', 'lhs_input_shape', _tuple_scalar_codec),
FieldEncoding('rhs_input_shape', 'rhs_input_shape', _tuple_scalar_codec),
FieldEncoding('in1_const_attrs', 'in1', _constant_attrs_codec),
FieldEncoding('in2_const_attrs', 'in2', optional_codec(_constant_attrs_codec), None)
)),
ConstructorEncoding(attributes.ExternalAttrs, 'external', _external_attrs_codec),
ConstructorEncoding(attributes.QNNQuantizeAttrs, 'qnn_quantize', (
FieldEncoding('out_dtype', 'out_dtype', scalar_codec),
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
FieldEncoding('output_scale', 'output_scale', ndarray_codec),
FieldEncoding('output_zero_point', 'output_zero_point', ndarray_codec)
)),
ConstructorEncoding(attributes.QNNDequantizeAttrs, 'qnn_dequantize', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
FieldEncoding('input_scale', 'input_scale', ndarray_codec),
FieldEncoding('input_zero_point', 'input_zero_point', ndarray_codec)
)),
ConstructorEncoding(attributes.RequantizeAttrs, 'qnn_requantize', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('rounding', 'rounding', scalar_codec),
FieldEncoding('compute_dtype', 'compute_dtype', scalar_codec),
FieldEncoding('out_dtype', 'out_dtype', scalar_codec),
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC)
)),
ConstructorEncoding(attributes.AwesomeAttributes, 'qnn_binary', ()),
ConstructorEncoding(attributes.MultiplyAttrs, 'multiply', _multiply_attrs_codec),
ConstructorEncoding(attributes.BatchMatmulAttrs, 'batch_matmul', (
FieldEncoding('transpose_b', 'transpose_b', scalar_codec),
FieldEncoding('input_shapes', 'input_shapes', _list_tuple_scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC)
)),
ConstructorEncoding(attributes.LayerNormAttrs, 'layer_norm', _layer_norm_attrs_codec),
ConstructorEncoding(attributes.SliceConcatAttrs, 'slice_concat', (
FieldEncoding('slice_attrs', 'slice_attrs', list_codec(_strided_slice_attrs_codec)),
FieldEncoding('tuple_concat_attrs', 'tuple_concat_attrs', _tuple_concat_attrs_codec),
)),
ConstructorEncoding(attributes.RMSNormAttrs, 'rms_norm', (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('epsilon', 'epsilon', scalar_codec),
FieldEncoding('scalar_type', 'scalar_type', SCALARTYPE_CODEC),
)),
ConstructorEncoding(attributes.InstanceNormAttrs, 'instance_norm', _instance_norm_attrs_codec),
ConstructorEncoding(attributes.GridSampleAttrs, 'grid_sample', _grid_sample_attrs_codec),
ConstructorEncoding(attributes.VarianceAttrs, 'variance', _variance_attrs_codec)
]
)
[docs]
AWESOMECALIBATTRS_CODEC = dataclass_style_codec(
attributes.AwesomeCalibAttrs,
(
FieldEncoding('observer', 'observer', null_value_codec, None),
FieldEncoding('intermediate_observers', 'intermediate_observers', null_value_codec, None),
FieldEncoding('quant', 'quant', data_value_codec(QUANT_RESULT_TENSOR_TYPE_CODEC)),
FieldEncoding('input_quant', 'input_quant', dict_codec(scalar_codec, data_value_codec(
QUANT_RESULT_TENSOR_TYPE_CODEC))),
FieldEncoding('precomputed_quant', 'precomputed_quant', optional_codec(QUANTIZATION_CODEC)),
)
)
[docs]
AWESOMEQUANTATTRS_CODEC = adt_style_codec(
'op',
[
ConstructorEncoding(attributes.PlaceholderQuantAttrs, 'placeholder', (
FieldEncoding('type', 'type', TENSORTYPE_CODEC),
FieldEncoding('quantization', 'quantization', optional_codec(QUANTIZATION_CODEC))
)),
ConstructorEncoding(attributes.ConstantQuantAttrs, 'constant', (
FieldEncoding('quant_data', 'data', ndarray_codec),
)),
ConstructorEncoding(attributes.PoolQuantAttrs, 'pool', (
FieldEncoding('pool_attrs', 'pool_attrs', AWESOMEATTRIBUTES_CODEC),
FieldEncoding('pad_value', 'pad_value', scalar_codec),
FieldEncoding('rounding_type', 'rounding_type', ROUNDTYPE_CODEC, RoundType.TOEVEN),
FieldEncoding('input_int16', 'input_int16', scalar_codec),
FieldEncoding('requant', 'requant', optional_codec(REQUANTIZATION_CODEC))
)),
ConstructorEncoding(attributes.AddQuantAttrs, 'add', (
FieldEncoding('lhs_input_shape', 'lhs_input_shape', _tuple_scalar_codec),
FieldEncoding('rhs_input_shape', 'rhs_input_shape', _tuple_scalar_codec),
FieldEncoding('input_int16', 'input_int16', scalar_codec),
FieldEncoding('requant', 'requant', REQUANTIZATION_CODEC),
FieldEncoding('relu_zero_point', 'relu_zero_point', scalar_codec, 0),
FieldEncoding('lhs_scale', 'lhs_scale', scalar_codec, 1),
FieldEncoding('rhs_scale', 'rhs_scale', scalar_codec, 1),
FieldEncoding('layer_bits', 'layer_bits', scalar_codec, 8),
FieldEncoding('activ_attrs', 'activ_attrs', optional_codec(_activ_quant_attrs_codec), None)
)),
ConstructorEncoding(attributes.MultiplyQuantAttrs, 'multiply', _multply_quant_attrs_codec),
ConstructorEncoding(attributes.ConcatQuantAttrs, 'concat', _concat_quant_attrs_codec),
ConstructorEncoding(attributes.ImageResize2DQuantAttrs, 'image_resize_2d', (
FieldEncoding('image_resize2d_attrs', 'image_resize2d_attrs', AWESOMEATTRIBUTES_CODEC),
FieldEncoding('input_zp', 'input_zp', scalar_codec, 0),
FieldEncoding('input_scale', 'input_scale', scalar_codec, 0),
FieldEncoding('rounding_type', 'rounding_type', ROUNDTYPE_CODEC, RoundType.TOEVEN)
)),
ConstructorEncoding(attributes.UDFQuantAttrs, 'udf', _udf_quant_attrs_codec),
ConstructorEncoding(attributes.DivideQuantAttrs, 'divide', (
FieldEncoding('udf_attrs', 'udf_attrs', _udf_quant_attrs_codec),
FieldEncoding('multiply_attrs', 'mulitply_attrs', _multply_quant_attrs_codec),
)),
ConstructorEncoding(attributes.LeakyReluCompositeQuantAttrs, 'leaky_relu', (
FieldEncoding('attrs', 'attrs', AWESOMEATTRIBUTES_CODEC),
FieldEncoding('leaky_relu_uses_udf', 'use_udf', scalar_codec, True),
FieldEncoding('leaky_relu_quant_attrs', 'leaky_relu', optional_codec(_leaky_relu_quant_attrs_codec), None),
FieldEncoding('udf_quant_attrs', 'udf', optional_codec(_udf_quant_attrs_codec), None)
)),
ConstructorEncoding(attributes.PReluQuantAttrs, 'prelu', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('quant_alpha', 'quant_alpha', ndarray_codec),
FieldEncoding('alpha_shift', 'alpha_shift', scalar_codec, 0),
FieldEncoding('data_zero_point', 'data_zero_point', scalar_codec, 0),
)),
ConstructorEncoding(attributes.ConvQuantAttrs, 'conv', (
FieldEncoding('conv_attrs', 'conv', _conv_attrs_codec),
FieldEncoding('weight_quant_data', 'weight_quant_data', ndarray_codec),
FieldEncoding('requant', 'requant', REQUANTIZATION_CODEC),
FieldEncoding('scale', 'scale', scalar_codec),
FieldEncoding('zero_point', 'zero_point', scalar_codec),
FieldEncoding('input_zp', 'input_zp', scalar_codec),
FieldEncoding('bias_quant_data', 'bias_quant_data', optional_codec(ndarray_codec)),
FieldEncoding('per_channel', 'per_channel', scalar_codec, attributes.DEFAULT_PER_CHANNEL),
FieldEncoding('activ_attrs', 'activ_attrs', optional_codec(_activ_quant_attrs_codec), None),
FieldEncoding('input_int16', 'input_int16', scalar_codec, False),
FieldEncoding('msb_left_shift', 'msb_left_shift', _scalar_or_ndarray_codec),
)),
ConstructorEncoding(attributes.MeanQuantAttrs, 'mean', (
FieldEncoding('attrs', 'attrs', AWESOMEATTRIBUTES_CODEC),
FieldEncoding('node_scales', 'node_scales', scalar_codec, 1.0),
FieldEncoding('node_zps', 'node_zps', scalar_codec, 0)
)),
ConstructorEncoding(attributes.LRNQuantAttrs, 'lrn', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('size', 'size', scalar_codec),
FieldEncoding('shape', 'shape', _tuple_scalar_codec),
FieldEncoding('input_zp', 'input_zp', scalar_codec),
FieldEncoding('lut_scale', 'lut_scale', scalar_codec),
FieldEncoding('lut_zp_corr', 'lut_zp_corr', scalar_codec),
FieldEncoding('lut_sh', 'lut_sh', scalar_codec),
FieldEncoding('output_scale', 'output_scale', scalar_codec),
FieldEncoding('output_zp_corr', 'output_zp_corr', scalar_codec),
FieldEncoding('output_sh', 'output_sh', scalar_codec),
FieldEncoding('lookup_table', 'lookup_table', optional_codec(ndarray_codec))
)),
ConstructorEncoding(attributes.SoftmaxQuantAttrs, 'softmax', (
FieldEncoding('axis', 'axis', scalar_codec),
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('exp_zp', 'exp_zp', scalar_codec),
FieldEncoding('rec_zp', 'rec_zp', scalar_codec),
FieldEncoding('requant_lut', 'requant_lut', REQUANTIZATION_CODEC),
FieldEncoding('requant_output', 'requant_output', REQUANTIZATION_CODEC),
FieldEncoding('lookup_table_exp', 'lookup_table_exp', optional_codec(ndarray_codec)),
FieldEncoding('lookup_table_rec', 'lookup_table_rec', optional_codec(ndarray_codec)),
FieldEncoding('enable_int16', 'enable_int16', scalar_codec),
FieldEncoding('lut_input_pre_shift', 'lut_input_pre_shift', optional_codec(scalar_codec)),
FieldEncoding('output_pre_shift', 'output_pre_shift', optional_codec(scalar_codec)),
)),
ConstructorEncoding(attributes.UpsamplingQuantAttrs, 'upsample', (
FieldEncoding('upsampling_attrs', 'attrs', _upsampling_attrs_codec),
FieldEncoding('input_zp', 'input_zp', scalar_codec, 0),
FieldEncoding('input_scale', 'input_scale', scalar_codec, 1.0),
FieldEncoding('rounding_type', 'rounding_type', ROUNDTYPE_CODEC, RoundType.TOEVEN),
FieldEncoding('input_int16', 'input_int16', scalar_codec)
)),
ConstructorEncoding(attributes.RequantizeQuantAttrs, 'qnn_requantize', (
FieldEncoding('attrs', 'attrs', AWESOMEATTRIBUTES_CODEC),
FieldEncoding('requant', 'requant', REQUANTIZATION_CODEC),
)),
ConstructorEncoding(attributes.SubtractQuantAttrs, 'qnn_subtract', (
FieldEncoding('attrs', 'attrs', AWESOMEATTRIBUTES_CODEC),
FieldEncoding('input_int16', 'input_int16', scalar_codec),
FieldEncoding('requant', 'requant', REQUANTIZATION_CODEC),
FieldEncoding('lhs_scale', 'lhs_scale', scalar_codec, 1),
FieldEncoding('rhs_scale', 'rhs_scale', scalar_codec, 1),
FieldEncoding('layer_bits', 'layer_bits', scalar_codec, 8),
)),
ConstructorEncoding(attributes.TupleAttrs, 'tuple', _tuple_attrs_codec),
ConstructorEncoding(attributes.TupleGetItemAttrs, 'tuple_get_item', (
FieldEncoding('input_types', 'input_types', _list_tensor_type_codec),
FieldEncoding('index', 'index', scalar_codec),
)),
ConstructorEncoding(attributes.LayoutTransformAttrs, 'layout_transform', (
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
FieldEncoding('src_layout', 'src_layout', scalar_codec),
FieldEncoding('dst_layout', 'dst_layout', scalar_codec),
FieldEncoding('implicitly_removable', 'implicitly_removable', scalar_codec, False)
)),
ConstructorEncoding(attributes.ArgMaxQuantAttrs, 'argmax', (
FieldEncoding('attrs', 'attrs', _argmax_attrs_codec),
)),
ConstructorEncoding(attributes.BatchMatmulQuantAttrs, 'batch_matmul', (
FieldEncoding('attrs', 'attrs', AWESOMEATTRIBUTES_CODEC),
FieldEncoding('lhs_zp', 'lhs_zp', scalar_codec),
FieldEncoding('rhs_zp', 'rhs_zp', scalar_codec),
FieldEncoding('requant', 'requant', REQUANTIZATION_CODEC),
FieldEncoding('intrinsic_shift', 'intrinsic_shift', scalar_codec),
)),
ConstructorEncoding(attributes.LayerNormQuantAttrs, 'layer_norm', _layer_norm_quant_attrs_codec),
ConstructorEncoding(attributes.SliceConcatQuantAttrs, 'slice_concat', (
FieldEncoding('slice_attrs', 'slice_attrs', list_codec(_strided_slice_attrs_codec)),
FieldEncoding('tuple_concat_attrs', 'tuple_concat_attrs', _concat_quant_attrs_codec),
)),
ConstructorEncoding(attributes.BroadcastToQuantAttrs, 'broadcast_to', (
FieldEncoding('input_type', 'input_type', TENSORTYPE_CODEC),
FieldEncoding('output_shape', 'output_shape', _tuple_scalar_codec),
)),
ConstructorEncoding(attributes.RMSNormQuantAttrs, 'rms_norm', (
FieldEncoding('input_shape', 'input_shape', _tuple_scalar_codec),
FieldEncoding('zp_ifm', 'zp_ifm', scalar_codec),
FieldEncoding('lookup_table_rsqrt', 'lookup_table_rsqrt', ndarray_codec),
FieldEncoding('zp_rsqrt', 'zp_rsqrt', scalar_codec),
FieldEncoding('requant_lut_input', 'requant_lut_input', REQUANTIZATION_CODEC),
FieldEncoding('requant_output', 'requant_output', REQUANTIZATION_CODEC),
FieldEncoding('lut_input_pre_shift', 'lut_input_pre_shift', scalar_codec),
FieldEncoding('output_pre_shift', 'output_pre_shift', scalar_codec),
FieldEncoding('enable_lut_int16', 'enable_lut_int16', scalar_codec),
)),
ConstructorEncoding(attributes.InstanceNormQuantAttrs, 'instance_norm', (
FieldEncoding('attrs', 'attrs', _instance_norm_attrs_codec),
FieldEncoding('lut_rsqrt', 'lut_rsqrt', ndarray_codec),
FieldEncoding('zp_rsqrt', 'zp_rsqrt', scalar_codec),
FieldEncoding('requant_out', 'requant_out', REQUANTIZATION_CODEC)
)),
ConstructorEncoding(attributes.VarianceQuantAttrs, 'variance', (
FieldEncoding('attrs', 'attrs', _variance_attrs_codec),
FieldEncoding('requant', 'requant', REQUANTIZATION_CODEC),
FieldEncoding('requant_var', 'requant_var', REQUANTIZATION_CODEC)
)),
ConstructorEncoding(attributes.ClipQuantAttrs, 'clip', _clip_attrs_codec),
ConstructorEncoding(attributes.ReluQuantAttrs, 'relu', _relu_quant_attrs_codec),
]
)
[docs]
SIMAIR_CODEC = dataclass_style_codec(
SiMaIR,
(
FieldEncoding('operation', 'op', AWESOMEOPERATION_CODEC),
FieldEncoding('_attrs', 'attrs', optional_codec(AWESOMEATTRIBUTES_CODEC), None),
FieldEncoding('calib_attrs', 'calib', optional_codec(AWESOMECALIBATTRS_CODEC)),
FieldEncoding('_quant_attrs', 'quant', optional_codec(AWESOMEQUANTATTRS_CODEC)),
FieldEncoding('quant_config', 'quant_config', optional_codec(QUANTIZATION_CONFIGS_CODEC)),
FieldEncoding('backend', 'backend', BACKEND_CODEC),
FieldEncoding('_metadata', 'metadata', SIMAIRMETADATA_CODEC)
)
)
# The IR in BackendIR is treated as an abstract data type by AFE.
# To support serialization, the external project that defines the IR
# should also implement serialization for it.
class _BackendIRCodec:
def encode(self, state: SerializerEncState, obj: backends.BackendIR) -> Any:
raise EncodeError("Serialization of backend code is not supported")
def decode(self, state: SerializerDecState, doc: Any) -> backends.BackendIR:
raise DecodeError("Deserialization of backend code is not supported")
[docs]
BACKENDIR_CODEC = _BackendIRCodec()
[docs]
AWESOMENODE_IR_CODEC = type_tagged_codec(
'node_type',
(
('operator', SiMaIR, SIMAIR_CODEC),
('subgraph', AwesomeNet, lazy_codec(lambda: AWESOMENET_CODEC)),
('backend', backends.BackendIR, BACKENDIR_CODEC)
)
)
[docs]
AWESOMENODE_CODEC = dataclass_style_codec(
AwesomeNode,
(
FieldEncoding('name', 'name', scalar_codec),
FieldEncoding('input_names', 'parameters', _list_scalar_codec),
FieldEncoding('input_node_names', 'inputs', _list_scalar_codec),
FieldEncoding('ir', 'ir', AWESOMENODE_IR_CODEC),
FieldEncoding('_status', 'status', STATUS_CODEC),
FieldEncoding('_layer_stats', 'layer_stats', optional_codec(LAYERSTATS_CODEC))
)
)
def _nodes_to_list(nodes: Dict[NodeName, AwesomeNode]) -> List[AwesomeNode]:
"""
Extract a list of nodes from a node mapping in AwesomeNet. The list will be serialized.
Assumes that nodes satisfies `nodes[x].name == x` for all x in nodes.keys().
:param nodes: Nodes taken from AwesomeNet
:return: List of nodes
"""
return list(nodes.values())
def _nodes_to_dict(nodes: List[AwesomeNode]) -> Dict[NodeName, AwesomeNode]:
"""
Create a node mapping for AwesomeNet that contains the given nodes.
:param nodes: Nodes to use
:return: Node mapping
"""
return {n.name: n for n in nodes}
[docs]
AWESOMENET_CODEC = dataclass_style_codec(
AwesomeNet,
(
FieldEncoding('name', 'name', scalar_codec),
FieldEncoding('nodes', 'nodes', map_codec(_nodes_to_dict, _nodes_to_list, list_codec(AWESOMENODE_CODEC))),
FieldEncoding('input_node_names', 'parameters', _list_scalar_codec),
FieldEncoding('output_node_name', 'output', scalar_codec),
FieldEncoding('_status', 'status', STATUS_CODEC),
FieldEncoding('_execution_order', 'execution_order', _list_scalar_codec),
FieldEncoding('_prune_dict', 'prune_dict', dict_codec(scalar_codec, _list_scalar_codec)),
FieldEncoding('_float_node_list', 'float_node_list', _list_scalar_codec),
FieldEncoding('_is_subgraph', 'is_subgraph', scalar_codec, default=False),
FieldEncoding('_backend', '_backend', BACKEND_CODEC, default=Backend.NONE),
FieldEncoding('_target', '_target', TARGET_CODEC)
)
)
[docs]
def encode_awesomenet(net: AwesomeNet) -> Tuple[Any, Dict[Any, np.ndarray]]:
"""
Encode an AwesomeNet in preparation for saving to a yaml file.
:param net: Data to encode
:return: A tuple of 2 items: a data structure that can be written to file by yaml's safe_dump, and
a map that can be written to file by numpy's savez.
"""
encoder_state = SerializerEncState()
yaml_data = AWESOMENET_CODEC.encode(encoder_state, net)
numpy_data = encoder_state.get_arrays()
return yaml_data, numpy_data
[docs]
def decode_awesomenet(yaml_data: Any, numpy_data: Dict[Any, np.ndarray]) -> AwesomeNet:
"""
Decode an AwesomeNet that was encoded by encode_awesomenet.
:param yaml_data: Data to decode
:param numpy_data: Numpy data to retrieve
:return: Decoded AwesomeNet
"""
decoder_state = SerializerDecState(numpy_dict=numpy_data)
return AWESOMENET_CODEC.decode(decoder_state, yaml_data)
def _yaml_headers(label: str, version: str) -> List[str]:
"""
Create headers for a YAML document.
"""
header1 = "This YAML document contains a serialized " + label
header2 = "Format version " + version
return [header1, header2]
[docs]
def wrap_yaml_document(yaml_data: Any, label: str, version: str, commit_id: Optional[str]) -> Any:
"""
Wraps YAML-encodable data with some sanity check information to make a complete yaml document.
:param yaml_data: Data to wrap.
:param label: Human-readable, short text description of the data type.
:param version: Version of the data format. This is used to detect writer/reader incompatibility.
:param commit_id: last Afe git commit ID.
:return: Wrapped data
"""
# If running from the source tree, the source tree's git commit ID should have been provided.
if not is_installed() and commit_id is None:
sima_logger.sima_log_warning('Cant get commit ID.')
return _yaml_headers(label, version) + [{'commit_id':commit_id}, yaml_data]
[docs]
def unwrap_yaml_document(yaml_data: Any, label: str, version: str, commit_id: Optional[str]) -> Any:
"""
Checks information in wrapped YAML and unwraps it. Raises a DecodeError if the data seems invalid.
:param yaml_data: Data to unwrap.
:param label: Human-readable, short text description of the data type.
:param version: Version of the data format. This is used to detect writer/reader incompatibility.
:param commit_id: last Afe git commit ID.
:return: Unwrapped data
"""
# Check values
headers_check = _yaml_headers(label, version)
commit_id_check = commit_id
# Check that data begins with header from wrap_yaml_document
if not isinstance(yaml_data, list) or len(yaml_data) < 2 or yaml_data[:-2] != headers_check:
raise DecodeError("Data does not seem like a valid YAML document")
# Unpacking yaml data
data = yaml_data[-1]
commit_id = yaml_data[-2]['commit_id']
# If running from the source tree, detect if the code has changed by comparing the git commit ID.
if not is_installed():
# Check if commit ID matches or if a commit ID is None
if commit_id != commit_id_check or commit_id is None:
sima_logger.sima_log_warning('Warning: file format may be incompatible. '
'Commit ID does not match or unverifiable because it is missing.')
return data
#########################################
# JSONEncoder to generate the Netron File
#########################################
from afe._tvm._defines import TVMGraphModule, TVMOp, TVMFunction
# A list of classes we can serialize to Json
[docs]
SERIALIZABLE_CLASSES = (
AwesomeNet,
AwesomeNode,
SiMaIR,
operations.AwesomeOperation,
attributes.AwesomeAttributes,
attributes.AwesomeCalibAttrs,
attributes.AwesomeQuantAttrBase,
TensorValue,
TupleValue,
Quantization,
QuantizationConfigs,
TensorType,
requantization.BaseRequantization,
requantization.Narrowing,
LayerStats
)
[docs]
class AwesomeEncoder(json.JSONEncoder):
"""
Can encode to JSON all classes in SERIALIZABLE_CLASSES
"""
def __init__(self, *args, **kwargs) -> None:
"""
:numpy_dict: A dictionary of keys to numpy arrays we populate when removing numpy arrays from JSON dict
:index: An integer used to create unique keys for the numpy dict
"""
[docs]
self.numpy_dict: Dict[str, np.ndarray] = dict()
super().__init__(*args, **kwargs)
[docs]
def default(self, obj: Any) -> Any:
"""
When encoding an object to json, if it belongs to one of the SERIALIZABLE_CLASSES we store it as a dict in json
along with a __decode_key__ we use to reconstruct the class later. If we encounter an numpy array, instead
of storing it in json we replace it with a string key and [key, the numpy array] pair in a dictionary.
"""
if isinstance(obj, SERIALIZABLE_CLASSES):
obj_dict = dict()
obj_dict.update(obj.__dict__)
for k, v in obj_dict.items():
if isinstance(v, tuple):
obj_dict[k] = {'vals': list(v), '__decode_key__': 'tuple'}
obj_dict['__decode_key__'] = type(obj).__name__
return obj_dict
elif type(obj) == dict:
return {'__decode_key__': 'dict', 'vals': list((k, self.default(v)) for k, v in obj.items())}
elif type(obj).__module__ == np.__name__:
arr_str = "np_arr_{}".format(self.index)
self.numpy_dict[arr_str] = obj
self.index += 1
return arr_str
elif type(obj) == ScalarType:
return {'__decode_key__': 'ScalarType', 'value': obj.value}
elif isinstance(obj, np.number):
# This typing covers all numeric scalar values defined by numpy
return {'__decode_key__': 'NumPy type', 'value': obj.__name__}
elif isinstance(obj, NodeAssociatedValue):
# NodeAssociatedValue is generic, doesn't work with default serialization
return {'__decode_key__': 'NodeAssociatedValue',
'inputs': self.default(obj.inputs),
'output': self.default(obj.output)}
elif isinstance(obj, (TVMGraphModule, TVMOp, TVMFunction)):
"""
External op edge case: ExternalAttributes contain graph runtime modules for use during network execution
To avoid saving these runtimes in extra .tar files we instead generate them from other attributes in
ExternalAttributes when the network is loaded.
"""
return None
# Let the base class default method raise the TypeError
return json.JSONEncoder.default(self, obj)