#########################################################
# Copyright (C) 2021 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: Joey Chou
#########################################################
import itertools
import numpy as np
from typing import Dict, Type, Union, Tuple, Optional, cast
from dataclasses import dataclass
from afe.backends import Backend
from afe.core.configs import QuantizationConfigs, RunConfigs
from afe.ir.defines import InputName, InputsQuantCast, map_data_value, DataValue, NodeReporter
import afe.ir.operations as afe_op
import afe.ir.attributes as afe_attrs
from afe.ir.quantization_interface import make_quantize_op_interface
from afe.ir.tensor_type import NodeType
[docs]
SiMaIRTensorTypes = Union[np.ndarray, Tuple[np.ndarray, ...]]
@dataclass
[docs]
class SiMaIR:
"""
:param operation: Contains the node's functions
:param _attrs: Floating-point operator attributes.
If the operator is quantized, it should be None.
:param calib_attrs: Calibration attributes
:param _quant_attrs: Quantized operator attributes.
If the operator is not quantized, it should be None.
:param quant_config: Parameters that influence how the operator is quantized.
This field is immutable.
:param backend: If this node is assigned to a backend by itself (not counting being
part of a subgraph assigned to a backend), this is the assigned backend.
Otherwise, this is Backend.NONE.
"""
[docs]
operation: afe_op.AwesomeOperation
_attrs: Optional[afe_attrs.AwesomeAttributes]
[docs]
calib_attrs: afe_attrs.AwesomeCalibAttrs
_quant_attrs: Optional[afe_attrs.AwesomeQuantAttrBase]
[docs]
quant_config: QuantizationConfigs
# Metadata
_metadata: SiMaIRMetadata
def _validate_attrs(self):
# SiMaIR must contain either attrs or quant_attrs
assert (self._attrs is None) != (self._quant_attrs is None)
def __post_init__(self):
self._validate_attrs()
@property
[docs]
def attrs(self) -> Optional[afe_attrs.AwesomeAttributes]:
self._validate_attrs()
return self._attrs
[docs]
def get_attrs(self) -> afe_attrs.AwesomeAttributes | afe_attrs.AwesomeQuantAttrBase:
self._validate_attrs()
if self._attrs is None:
return self._quant_attrs
else:
return self._attrs
@attrs.setter
def attrs(self, _attrs: Optional[afe_attrs.AwesomeAttributes]):
self._attrs = _attrs
self._quant_attrs = None
@property
[docs]
def quant_attrs(self) -> Optional[afe_attrs.AwesomeQuantAttrBase]:
self._validate_attrs()
return self._quant_attrs
@quant_attrs.setter
def quant_attrs(self, _quant_attrs: Optional[afe_attrs.AwesomeQuantAttrBase]):
self._quant_attrs = _quant_attrs
self._attrs = None
[docs]
def get_type(self) -> NodeType:
"""
Get the type of this node.
"""
a = self.quant_attrs if self.quant_attrs is not None else self.attrs
return self.operation.get_type(a)
[docs]
def run(self, inputs: SiMaIRInputTypes, config: RunConfigs):
return self.operation.run(self.attrs, inputs, config)
[docs]
def calibrate(self, inputs: SiMaIRInputTypes, config: RunConfigs):
return self.operation.calibrate(self.attrs, self.calib_attrs, inputs, config)
[docs]
def quantize(self,
inputs: Dict[InputName, afe_op.QuantizationTensorData],
placeholder_value: Optional[afe_op.QuantizationTensorData],
quantize_attributes: bool,
error_reporter: Optional[NodeReporter] = None) \
-> Tuple[afe_op.QuantizationTensorData, InputsQuantCast]:
"""
Select quantization scales of input and output tensors, and quantize this operation.
The input and output attributes are planned to be removed, and they should not be used in new code.
:param inputs: Properties of the inputs.
It has quantization scales of the input tensors and attributes of the nodes that calculate the inputs.
:param placeholder_value: If the node being quantized is a placeholder node, the properties of
the placeholder's input data. None otherwise.
:param quantize_attributes: If True, the operation will be quantized. If False, the operation will not be
quantized. Either way, quantization scales will be selected.
:param error_reporter: Node reporter of the node to be quantized.
:return: Properties of the node's output and casts that should be applied to the node's input.
It has quantization scales of the output and attributes of this node.
"""
quant_inputs = {k : v[:2] for k, v in inputs.items()}
calibration_inputs = {k : v[2] for k, v in inputs.items()}
output_distribution, intermediate_distributions = (
self.operation.get_observed_distribution(self.calib_attrs, quant_inputs))
if placeholder_value is not None:
assert isinstance(self.operation, afe_op.PlaceholderOp)
placeholder_quantization = placeholder_value[0]
else:
assert not isinstance(self.operation, afe_op.PlaceholderOp)
placeholder_quantization = None
if quantize_attributes:
quant_interface, quant_result = \
make_quantize_op_interface(inputs, placeholder_quantization, output_distribution,
intermediate_distributions)
quant_attrs = self.operation.quantize(self.attrs, quant_interface, self.quant_config,
error_reporter)
provided_input_types = {n: v[0] for n, v in inputs.items()}
wanted_input_types = quant_result.get_result().inputs
result_type = quant_result.get_result().output
# Record the type of the quantized node
self.calib_attrs.input_quant = wanted_input_types
self.calib_attrs.quant = result_type
# Cast all inputs that need it
quantization_casts = afe_op.make_quantization_casts(provided_input_types, wanted_input_types)
# Save the created quantization attributes and reset original attributes to None.
if isinstance(quant_attrs, afe_attrs.AwesomeQuantAttrBase):
self.quant_attrs = quant_attrs
elif isinstance(quant_attrs, afe_attrs.AwesomeAttributes):
self.attrs = quant_attrs
return (result_type, output_distribution), quantization_casts
else:
# # Do not quantize. Compute quantization scales to allow other code to be quantized.
# Cast inputs to match the operation's type.
# Can get the type with get_type since it did not change.
node_type = self.operation.get_type(self.attrs)
wanted_types = {cast(InputName, k): map_data_value(afe_attrs.QuantResultTensorType.from_type, tt)
for k, tt in node_type.inputs.items()}
result_type = map_data_value(afe_attrs.QuantResultTensorType.from_type, node_type.output)
provided_types = {k: v[0] for k, v in inputs.items()}
input_cast = afe_op.make_quantization_casts(provided_types, wanted_types)
# Record that the node is not quantized
self.calib_attrs.input_quant = wanted_types
self.calib_attrs.quant = result_type
return (result_type, output_distribution), input_cast
[docs]
def run_quant(self, inputs: SiMaIRInputTypes, config: RunConfigs):
return self.operation.run_quant(self.quant_attrs, inputs, config)
[docs]
def set_batch_size(self, batch_size: int):
"""
Modifies SiMaIR's internal parameters to accommodate for a given batch size.
:param batch_size: Integer value representing the batch size of the inputs to the AwesomeNet.
"""
self.calib_attrs.set_batch_size(batch_size)
if self.attrs is not None:
self.attrs.set_batch_size(batch_size)
if self.quant_attrs is not None:
self.quant_attrs.set_batch_size(batch_size)
###############################
# Used by the node constructor
###############################
@dataclass(frozen=True)
[docs]
class SiMaIRParamsDict:
[docs]
attrs: Type[afe_attrs.AwesomeAttributes]
[docs]
quant_attrs: Type[afe_attrs.AwesomeQuantAttrBase]
[docs]
operation: afe_op.AwesomeOperation
# Translations from TVM operators to SiMa IR operators, for those TVM operators that translate to a single
# SiMa IR operator and that have the same attributes in TVM as in SiMa IR.
# This table will be removed. A table will be created in _operators.py to replace it.
[docs]
NODE_DICT_OPERATORS: Dict[str, SiMaIRParamsDict] = {
# SINGLE ATTRIBUTES
"placeholder": SiMaIRParamsDict(afe_attrs.AwesomeAttributes, afe_attrs.AwesomeQuantAttrBase, afe_op.PlaceholderOp()),
"constant": SiMaIRParamsDict(afe_attrs.ConstantAttrs, afe_attrs.ConstantQuantAttrs, afe_op.ConstantOp()),
"nn.adaptive_avg_pool2d": SiMaIRParamsDict(afe_attrs.AdaptiveAvgPool2DAttrs, afe_attrs.PoolQuantAttrs, afe_op.AdaptiveAvgPool2DOp()),
"multiply": SiMaIRParamsDict(afe_attrs.AwesomeAttributes, afe_attrs.MultiplyQuantAttrs, afe_op.MultiplyOp()),
"tuple": SiMaIRParamsDict(afe_attrs.TupleAttrs, afe_attrs.AwesomeQuantAttrBase, afe_op.TupleOp()),
"nn.relu": SiMaIRParamsDict(afe_attrs.ReluAttrs, afe_attrs.AwesomeQuantAttrBase, afe_op.ReluOp()),
}
# Translations from TVM operators to SiMa IR operators, for operators that are handled specially
# in the TVM converter.
# This table table will be removed.
[docs]
NODE_DICT_SPECIAL: Dict[str, SiMaIRParamsDict] = {
# CUSTOM OPERATIONS
"nn.custom_op_1": SiMaIRParamsDict(afe_attrs.CustomOpAttrs, afe_attrs.AwesomeQuantAttrBase, afe_op.CustomOp()),
"nn.custom_op_2": SiMaIRParamsDict(afe_attrs.CustomOpAttrs, afe_attrs.AwesomeQuantAttrBase, afe_op.CustomOp()),
"nn.custom_op_3": SiMaIRParamsDict(afe_attrs.CustomOpAttrs, afe_attrs.AwesomeQuantAttrBase, afe_op.CustomOp()),
"nn.custom_op_4": SiMaIRParamsDict(afe_attrs.CustomOpAttrs, afe_attrs.AwesomeQuantAttrBase, afe_op.CustomOp()),
"nn.custom_op_5": SiMaIRParamsDict(afe_attrs.CustomOpAttrs, afe_attrs.AwesomeQuantAttrBase, afe_op.CustomOp()),
# TEMPORARY HACK
"constant_multiply_add": SiMaIRParamsDict(afe_attrs.ConstantMultiplyAddAttrs, afe_attrs.AddQuantAttrs, afe_op.ConstantMultiplyAddOp())
}
# Dictionary of tvm expression names to afe_attrs.AwesomeAttributes, afe_attrs.AwesomeQuantAttr, afe_op.AwesomeOperation.
[docs]
NODE_DICT: Dict[str, SiMaIRParamsDict] = \
dict(itertools.chain(NODE_DICT_OPERATORS.items(), NODE_DICT_SPECIAL.items()))
###############################
# Used by the Serializer
###############################
# Dictionary of afe_attrs.AwesomeAttributes class names to afe_attrs.AwesomeAttributes.
[docs]
ATTRS_NAME_DICT: Dict[str, Type[afe_attrs.AwesomeAttributes]] = {
# SINGLE ATTRIBUTES
"AwesomeAttributes": afe_attrs.AwesomeAttributes,
"PlaceholderAttrs": afe_attrs.PlaceholderAttrs,
"ConstantAttrs": afe_attrs.ConstantAttrs,
"MaxPoolAttrs": afe_attrs.MaxPoolAttrs,
"AvgPoolAttrs": afe_attrs.AvgPoolAttrs,
"AdaptiveAvgPool2DAttrs": afe_attrs.AdaptiveAvgPool2DAttrs,
"ReluAttrs": afe_attrs.ReluAttrs,
"AddAttrs": afe_attrs.AddAttrs,
"BiasAddAttrs": afe_attrs.BiasAddAttrs,
"MeanAttrs": afe_attrs.MeanAttrs,
"SqueezeAttrs": afe_attrs.SqueezeAttrs,
"ArgMaxAttrs": afe_attrs.ArgMaxAttrs,
"SoftmaxAttrs": afe_attrs.SoftmaxAttrs,
"TupleAttrs": afe_attrs.TupleAttrs,
"PadAttrs": afe_attrs.PadAttrs,
"LRNAttrs": afe_attrs.LRNAttrs,
"ConcatenateAttrs": afe_attrs.ConcatenateAttrs,
"TransposeAttrs": afe_attrs.TransposeAttrs,
"ClipAttrs": afe_attrs.ClipAttrs,
"ReshapeAttrs": afe_attrs.ReshapeAttrs,
"ExpandDimsAttrs": afe_attrs.ExpandDimsAttrs,
"BatchFlattenAttrs": afe_attrs.BatchFlattenAttrs,
"UpsamplingAttrs": afe_attrs.UpsamplingAttrs,
"LayoutTransformAttrs": afe_attrs.LayoutTransformAttrs,
"TessellationTransformAttrs": afe_attrs.TessellationTransformAttrs,
"DetessellationTransformAttrs": afe_attrs.DetessellationTransformAttrs,
"PackTransformAttrs": afe_attrs.PackTransformAttrs,
"UnpackTransformAttrs": afe_attrs.UnpackTransformAttrs,
"NormalizationTransformAttrs": afe_attrs.NormalizationTransformAttrs,
"QuantizationTransformAttrs": afe_attrs.QuantizationTransformAttrs,
"DequantizationTransformAttrs": afe_attrs.DequantizationTransformAttrs,
"ResizeTransformAttrs": afe_attrs.ResizeTransformAttrs,
"ChromaUpsampleTransformAttrs": afe_attrs.ChromaUpsampleTransformAttrs,
"YuvRgbConversionTransformAttrs": afe_attrs.YuvRgbConversionTransformAttrs,
"BgrRgbConversionTransformAttrs": afe_attrs.BgrRgbConversionTransformAttrs,
"SigmoidTransformAttrs": afe_attrs.SigmoidTransformAttrs,
"NmsMaxpoolTransformAttrs": afe_attrs.NmsMaxpoolTransformAttrs,
"CastAttrs": afe_attrs.CastAttrs,
"SubtractAttrs": afe_attrs.SubtractAttrs,
"ExtmAttrs": afe_attrs.ExtmAttrs,
"SumAttrs": afe_attrs.SumAttrs,
"ProdAttrs": afe_attrs.ProdAttrs,
"TupleGetItemAttrs": afe_attrs.TupleGetItemAttrs,
"MaximumAttrs": afe_attrs.MaximumAttrs,
"MinimumAttrs": afe_attrs.MinimumAttrs,
"FullAttrs": afe_attrs.FullAttrs,
"TileAttrs": afe_attrs.TileAttrs,
"SplitAttrs": afe_attrs.SplitAttrs,
"TakeAttrs": afe_attrs.TakeAttrs,
"StridedSliceAttrs": afe_attrs.StridedSliceAttrs,
"ImageResize2DAttrs": afe_attrs.ImageResize2DAttrs,
# UDF ATTRIBUTES
"UDFAttrs": afe_attrs.UDFAttrs,
"LeakyReluAttrs": afe_attrs.LeakyReluAttrs,
# COMPOSITE ATTRIBUTES
"ConvAddActivationAttrs": afe_attrs.ConvAddActivationAttrs,
"AddActivationAttrs": afe_attrs.AddActivationAttrs,
"TupleConcatenateAttrs": afe_attrs.TupleConcatenateAttrs,
"SwishAttrs": afe_attrs.SwishAttrs,
"PReluAttrs": afe_attrs.PReluAttrs,
"ConstantMultiplyAddAttrs": afe_attrs.ConstantMultiplyAddAttrs,
# PARTITIONING ATTRIBUTES
"ExternalAttrs": afe_attrs.ExternalAttrs,
# QNN Attributes
"QNNQuantizeAttrs": afe_attrs.QNNQuantizeAttrs,
"RequantizeAttrs": afe_attrs.RequantizeAttrs,
"QNNDequantizeAttrs": afe_attrs.QNNDequantizeAttrs,
}