#########################################################
# 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
#########################################################
"""
Operator-specific conversion from SiMaIR operators to
plugin sections in a MPK JSON file. An operator can be
converted to MPK JSON data by calling make_ev_plugin_parameters.
The result should be combined with operator-independent data
to make a PluginMPKData.
"""
from dataclasses import dataclass
from typing import List, Callable, Union, Mapping, Type, Dict, Any, Tuple, Optional
from afe.backends.mpk.defines import (
ConfigParamsMPKData, PluginInputNodeMPKData, InOutNodesMPKData, EV74ConfigParamsMPKData
)
from afe.backends.mpk.node import get_node_size, get_output_node_names
from afe.ir import attributes as afe_attrs, operations as afe_op
from afe.ir.defines import TensorValue, NodeName, data_value_elements
from afe.ir.node import AwesomeNode
from afe.ir.operations import get_strided_slice_out_shape
from afe.ir.sima_ir import SiMaIR
from afe.ir.tensor_type import scalar_type_to_dtype, data_byte_size
from sima_utils.logging import sima_logger
@dataclass
[docs]
class PluginMPKOperatorData:
"""
Fields of PluginMPKData that are produced in an operator-specific way.
These fields should be copied into a PluginMPKData to make a plugin.
"""
[docs]
config_params: ConfigParamsMPKData
[docs]
output_nodes: List[InOutNodesMPKData]
class _AfeToMPKOperationConverter:
"""
Extractor of MPK plugin parameters from an AwesomeNode.
Each class instance deals with a single operator.
"""
@classmethod
def gen_plugin_params(cls, node: AwesomeNode, input_node_names: List[NodeName],
desired_batch_size: int, actual_batch_size: int,
output_names: Dict[str, str]) -> PluginMPKOperatorData:
"""
Create operator-specific parameters for a plugin invocation corresponding
to the given node's operation.
"""
raise NotImplementedError("gen_plugin_params is abstract")
def _ev_attrs_to_params(kernel_name: str,
f: Callable[[Union[afe_attrs.AwesomeAttributes, afe_attrs.AwesomeQuantAttrBase]],
Dict[str, Any]]) -> _AfeToMPKOperationConverter:
"""
Construct an AfeToEVOperationConverter for an EV operator using
the given function to construct plugin parameters.
The input_nodes and output_nodes are converted from the AwesomeNode's
input and output.
:param kernel_name: The name of the EV kernel for the EV operator.
:param f: How to create kernel parameters for ConfigParamsMPKData
from data in an AwesomeNode
:return: Converter that produces PluginMPKOperatorData
"""
class Converter(_AfeToMPKOperationConverter):
@staticmethod
def _get_node_shapes_dict(node: AwesomeNode) -> Dict[str, Any]:
node_types = node.get_type()
input_shapes = []
for node_input_type in node_types.inputs.values():
input_shapes.extend([tt.shape for tt in data_value_elements(node_input_type)])
output_shapes = [tt.shape for tt in data_value_elements(node_types.output)]
return {"input_shapes": input_shapes, "output_shapes": output_shapes}
@classmethod
def gen_plugin_params(cls, node: AwesomeNode, input_nodes: List[PluginInputNodeMPKData],
desired_batch_size: int, actual_batch_size: int,
output_names: Dict[str, str]) -> PluginMPKOperatorData:
assert isinstance(node.ir, SiMaIR)
processor = "EV74"
attrs = node.ir.attrs
config = EV74ConfigParamsMPKData(desired_batch_size=desired_batch_size,
actual_batch_size=actual_batch_size,
kernel=kernel_name,
params=f(attrs) | cls._get_node_shapes_dict(node))
# Handle Unpack transform separately, as it generates multiple outputs.
if isinstance(attrs, afe_attrs.UnpackTransformAttrs):
output_nodes = [InOutNodesMPKData(name=f"{node.name}_{idx}",
size=data_byte_size(TensorValue(output_type)))
for idx, output_type in enumerate(attrs.tensor_types)]
else:
out_node_name = get_output_node_names(node.name, node.get_type().output, output_names)
assert isinstance(out_node_name, TensorValue)
output_nodes = [InOutNodesMPKData(name=out_node_name.value, size=get_node_size(node))]
return PluginMPKOperatorData(processor, config, input_nodes, output_nodes)
return Converter()
def _convert_argmax_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.ArgMaxAttrs)
assert attrs.result_scalar_type == afe_op.ScalarType.int32, \
"Argmax operator assigned to EV must have int32 output type"
assert len(attrs.axis) == 1, "Argmax operator assigned to EV must have a single axis"
axis, = attrs.axis
return {'axis': axis}
def _convert_layout_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.LayoutTransformAttrs)
return {'src_layout': attrs.src_layout, 'dst_layout': attrs.dst_layout}
def _convert_tessellation_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.TessellationTransformAttrs)
return {
'slice_shape': attrs.slice_shape,
'align_c16': attrs.align_c16,
'cblock': attrs.cblock,
'frame_type': scalar_type_to_dtype(attrs.frame_type.scalar)
}
def _convert_detessellation_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.DetessellationTransformAttrs)
return {
'slice_shape': attrs.slice_shape,
'align_c16': attrs.align_c16,
'cblock': attrs.cblock,
'frame_type': scalar_type_to_dtype(attrs.frame_type.scalar),
'frame_shape': list(attrs.frame_type.shape)
}
def _convert_pack_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.PackTransformAttrs)
return {}
def _convert_unpack_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.UnpackTransformAttrs)
return {
'tensor_types': [scalar_type_to_dtype(tensor_type.scalar) for tensor_type in attrs.tensor_types],
'tensor_shapes': [list(tensor_type.shape) for tensor_type in attrs.tensor_types]
}
def _convert_normalization_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.NormalizationTransformAttrs)
return {'channel_params': attrs.channel_params}
def _convert_quantization_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.QuantizationTransformAttrs)
return {
'channel_params': attrs.channel_params,
'num_bits': attrs.num_bits,
'rounding': attrs.rounding.name,
'output_data_type': scalar_type_to_dtype(attrs.output_data_type)
}
def _convert_dequantization_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.DequantizationTransformAttrs)
data_type = scalar_type_to_dtype(attrs.input_type.scalar)
return {'channel_params': attrs.channel_params,
'input_data_type': data_type}
def _convert_resize_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.ResizeTransformAttrs)
return {
'target_height': attrs.target_height,
'target_width': attrs.target_width,
'keep_aspect': attrs.keep_aspect,
'deposit_location': attrs.deposit_location.value,
'method': attrs.method.value
}
def _convert_chroma_upsample_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.ChromaUpsampleTransformAttrs)
return {
'frame_height': attrs.frame_height,
'frame_width': attrs.frame_width,
'yuv_sampling': attrs.yuv_sampling.value
}
def _convert_yuv_rgb_conversion_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.YuvRgbConversionTransformAttrs)
return {
'conversion': attrs.conversion.value,
'std': attrs.std.value
}
def _convert_bgr_rgb_conversion_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.BgrRgbConversionTransformAttrs)
return {
'conversion': attrs.conversion.value
}
def _convert_slice_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.StridedSliceAttrs)
output_shape = get_strided_slice_out_shape(attrs)
return {'begin': attrs.begin,
'end': attrs.end,
'input_shape': attrs.input_shape,
'output_shape': output_shape}
def _convert_sigmoid_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.SigmoidTransformAttrs)
return {'save_int16': attrs.save_int16}
def _convert_nms_maxpool_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.NmsMaxpoolTransformAttrs)
return {'kernel': attrs.kernel}
def _convert_softmax_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.SoftmaxAttrs)
return {'axis': attrs.axis}
def _convert_reshape_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.ReshapeAttrs)
return {'newshape': attrs.newshape}
def _convert_cast_transform_params(attrs: Union[afe_attrs.AwesomeAttributes,
afe_attrs.AwesomeQuantAttrBase]) -> Dict[str, Any]:
assert isinstance(attrs, afe_attrs.CastAttrs)
return {'out_dtype': attrs.out_dtype}
_OPERATOR_CONVERTERS: Mapping[Type[afe_op.AwesomeOperation], _AfeToMPKOperationConverter] = {
afe_op.ArgMaxOp: _ev_attrs_to_params("argmax", _convert_argmax_params),
afe_op.LayoutTransformOp: _ev_attrs_to_params("layout_transform", _convert_layout_transform_params),
afe_op.TessellationTransformOp: _ev_attrs_to_params("tessellation_transform",
_convert_tessellation_transform_params),
afe_op.DetessellationTransformOp: _ev_attrs_to_params("detessellation_transform",
_convert_detessellation_transform_params),
afe_op.PackTransformOp: _ev_attrs_to_params("pack_transform", _convert_pack_transform_params),
afe_op.UnpackTransformOp: _ev_attrs_to_params("unpack_transform", _convert_unpack_transform_params),
afe_op.NormalizationTransformOp: _ev_attrs_to_params("normalization_transform",
_convert_normalization_transform_params),
afe_op.QuantizationTransformOp: _ev_attrs_to_params("quantization_transform",
_convert_quantization_transform_params),
afe_op.DequantizationTransformOp: _ev_attrs_to_params("dequantization_transform",
_convert_dequantization_transform_params),
afe_op.ResizeTransformOp: _ev_attrs_to_params("resize_transform",
_convert_resize_transform_params),
afe_op.ChromaUpsampleTransformOp: _ev_attrs_to_params("chroma_upsample_transform",
_convert_chroma_upsample_transform_params),
afe_op.YuvRgbConversionTransformOp: _ev_attrs_to_params("yuv_rgb_conversion_transform",
_convert_yuv_rgb_conversion_transform_params),
afe_op.BgrRgbConversionTransformOp: _ev_attrs_to_params("bgr_rgb_conversion_transform",
_convert_bgr_rgb_conversion_transform_params),
afe_op.StridedSliceOp: _ev_attrs_to_params("slice_transform", _convert_slice_transform_params),
afe_op.SigmoidTransformOp: _ev_attrs_to_params("sigmoid_transform", _convert_sigmoid_transform_params),
afe_op.NmsMaxpoolTransformOp: _ev_attrs_to_params("nms_maxpool_transform", _convert_nms_maxpool_transform_params),
afe_op.SoftmaxOp: _ev_attrs_to_params("softmax_transform", _convert_softmax_transform_params),
afe_op.ReshapeOp: _ev_attrs_to_params("reshape_transform", _convert_reshape_transform_params),
afe_op.CastOp: _ev_attrs_to_params("cast_transform", _convert_cast_transform_params),
}
[docs]
def make_ev_plugin_parameters(node: AwesomeNode, input_nodes: List[PluginInputNodeMPKData],
desired_batch_size: int, actual_batch_size: int,
output_names: Dict[str, str]) -> PluginMPKOperatorData:
"""
Do the operator-specific part of extracting a node's parameters for an EV plugin.
"""
assert isinstance(node.ir, SiMaIR)
# SWMLA-5451: graceful error out when the same model input is needed for both int8 and int16.
if isinstance(node.ir.operation, afe_op.RequantizeOp):
raise sima_logger.UserFacingException(
"Unsupported int16 quantization use case where RequantizeOp is outside subgraph. "
"This happens when an input is quantized in int16 precision, but the operator "
"does not support int16 input. Please use int8 quantization instead."
)
try:
converter = _OPERATOR_CONVERTERS[type(node.ir.operation)]
except KeyError:
attrs = node.ir.attrs
operator_description = attrs.operations if isinstance(attrs, afe_attrs.ExternalAttrs) \
else node.ir.operation.__class__.__name__
raise TypeError(f"Operator(s) not supported on EV backend: {operator_description}")
return converter.gen_plugin_params(node, input_nodes, desired_batch_size, actual_batch_size, output_names)