#########################################################
# 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
#########################################################
"""
Functions for building and initializing AwesomeNodes
for specific operators. These functions provide a simpler
way to create an instance of a chosen operator and initialize
all the data structure fields. These functions are not
used for transformations on arbitrary operators.
"""
from typing import List, Tuple, Optional
import numpy as np
from afe.backends import Backend
from afe.core.configs import QuantizationConfigs
from afe.ir import attributes as attributes, operations as operations
from afe.ir.defines import NodeName, QuantCast, InputName, Status, DequantCast, TensorValue, Quantization, \
TupleValue, data_value_elements, get_expected_tensor_value, RequantMethod
from afe.ir.node import AwesomeNode
from afe.ir.sima_ir import SiMaIR, SiMaIRMetadata
from afe.ir.tensor_type import TensorType, ScalarType
from ml_kernels.math_helpers import RoundType
from ml_kernels.requantization import BaseRequantization
def _create_sima_ir(operation: operations.AwesomeOperation,
attrs: attributes.AwesomeAttributes | attributes.AwesomeQuantAttrBase,
calib_attrs: attributes.AwesomeCalibAttrs,
backend: Backend) -> SiMaIR:
"""
Creates a SiMaIR dataclass for the given operation.
:param operation: Operation for which the SiMaIR should be created.
:param attrs: SiMaIR operation attributes.
:param calib_attrs: SiMaIR operation calibration attributes.
:param backend: Backend on which the operator shall be executed.
:return: Operation's SiMaIR dataclass instance.
"""
_attrs = attrs if isinstance(attrs, attributes.AwesomeAttributes) else None
_quant_attrs = attrs if isinstance(attrs, attributes.AwesomeQuantAttrBase) else None
return SiMaIR(operation=operation, _attrs=_attrs, calib_attrs=calib_attrs, _quant_attrs=_quant_attrs,
quant_config=QuantizationConfigs(), backend=backend, _metadata=SiMaIRMetadata(""))
[docs]
def create_tuple_output_node(tuple_inputs: List[AwesomeNode], prefix: str) -> AwesomeNode:
"""
Creates a Tuple AwesomeNode. Used to accommodate AwesomeNet with multiple outputs.
:param tuple_inputs: List of AwesomeNodes that serve as input nodes to resulting
Tuple AwesomeNode.
:return: The created Tuple AwesomeNode.
"""
attrs = attributes.TupleAttrs([get_expected_tensor_value(t_i.get_type().output) for t_i in tuple_inputs])
calib_attrs = attributes.AwesomeCalibAttrs(
input_quant={f"input_{idx}": t_i.ir.calib_attrs.quant for idx, t_i in enumerate(tuple_inputs)},
quant=TupleValue([t_i.ir.calib_attrs.quant for t_i in tuple_inputs])
)
return AwesomeNode(name=f"{prefix}/tuple_out",
input_names=[f"input_{idx}" for idx in range(len(tuple_inputs))],
input_node_names=[t_i.name for t_i in tuple_inputs],
ir=_create_sima_ir(operations.TupleOp(), attrs, calib_attrs, Backend.NONE))
[docs]
def create_tuple_get_item_nodes(input_node_name: NodeName,
output_types: List[TensorType],
prefix: str,
input_quant: Optional[TupleValue[attributes.QuantResultTensorType]] = None) \
-> List[AwesomeNode]:
"""
Creates a TupleGetItem node that is used to split the Tuple output node.
:param input_node_name: The name of the input Tuple node.
:param output_types: List of types corresponding to the output types of the Tuple node.
:param prefix: The prefix string that is prepended to the TupleGetItem node name.
:param input_quant: Quantization parameters of the input nodes.
:return: The list of TupleGtItem AwesomeNodes.
"""
tgi_nodes: List[AwesomeNode] = list()
for idx in range(len(output_types)):
attrs = attributes.TupleGetItemAttrs(output_types, idx)
calib_attrs = attributes.AwesomeCalibAttrs() if not input_quant \
else attributes.AwesomeCalibAttrs(quant=TensorValue(data_value_elements(input_quant)[idx]),
input_quant={'tuple_value': input_quant})
node = AwesomeNode(name=f"{prefix}/tuple_get_item_{idx}",
input_names=["tuple_value"],
input_node_names=[input_node_name],
ir=_create_sima_ir(operations.TupleGetItemOp(), attrs, calib_attrs, Backend.NONE))
tgi_nodes.append(node)
return tgi_nodes
[docs]
def create_placeholder_node(node_name: str, input_type: TensorType) -> AwesomeNode:
"""
Creates a Placeholder AwesomeNode. Each AwesomeNet must have PlaceholderNodes to
hold the input data.
:param node_name: The name of the placeholder node to be created.
:param input_type: The input type for the placeholder node.
:return: Placeholder node that is created.
"""
attrs = attributes.PlaceholderAttrs(input_type)
calib_attrs = attributes.AwesomeCalibAttrs()
return AwesomeNode(name=node_name, input_names=['data'], input_node_names=[node_name],
ir=_create_sima_ir(operations.PlaceholderOp(), attrs, calib_attrs, Backend.NONE))
def _create_channel_params(scale: float, zero_point: int, num_channels: int) -> List[Tuple[float, int]]:
return [(scale, zero_point) for _ in range(num_channels)]
[docs]
def create_quantization_node(
input_name: NodeName, name_counter: int, cast: QuantCast, backend: Backend = Backend.EV
) -> AwesomeNode:
"""
Create a quantization node to quantize from float to integer (int8, int16)
:param input_name: Input node name
:param name_counter: Node index
:param cast: Cast to perform quantization
:param backend: Backend on which node will be executed
:return: AwesomeNode
"""
quant_type = cast.out_type
attrs = attributes.QuantizationTransformAttrs(
channel_params=_create_channel_params(cast.scale, cast.zero_point, 1),
input_shape=cast.shape, num_bits=cast.num_bits, output_data_type=quant_type
)
quant_result_tensor_type = \
attributes.QuantResultTensorType(
type=TensorType(scalar=quant_type, shape=cast.shape),
quant=Quantization.representable(cast.scale, cast.zero_point, cast.num_bits),
requant_method=RequantMethod.fractional_zero
)
input_quant_result_tensor_type = attributes.QuantResultTensorType(
type=TensorType(scalar=ScalarType.float32, shape=cast.shape), quant=None,
requant_method=None)
calib_attrs = attributes.AwesomeCalibAttrs(
quant=TensorValue(quant_result_tensor_type),
input_quant={InputName("data"): TensorValue(input_quant_result_tensor_type)}
)
return AwesomeNode(name=NodeName(f"quantize_{name_counter}"),
input_names=[InputName("data")],
input_node_names=[input_name],
ir=_create_sima_ir(
operations.QuantizationTransformOp(), attrs, calib_attrs, backend
),
_status=Status.SIMA_QUANTIZED)
[docs]
def create_dequantization_node(
input_name: NodeName, name_counter: int, cast: DequantCast, backend: Backend = Backend.EV
) -> AwesomeNode:
"""
Create a dequantization node to dequantize from integer (int8, int16, int32) to float.
:param input_name: Input node name
:param name_counter: Node index
:param cast: Cast to perform dequantization
:param backend: Backend on which node will be executed
:return: AwesomeNode
"""
input_scalar_type = ScalarType.from_numpy(cast.input_dtype)
output_scalar_type = ScalarType.from_numpy(cast.output_dtype)
attrs = attributes.DequantizationTransformAttrs(
channel_params=_create_channel_params(cast.scale, cast.zero_point, 1),
input_type=TensorType(input_scalar_type, cast.shape), output_type=output_scalar_type
)
quant_result_tensor_type = \
attributes.QuantResultTensorType(type=TensorType(scalar=output_scalar_type, shape=cast.shape),
quant=Quantization.representable(cast.scale, cast.zero_point, 8),
requant_method=RequantMethod.fractional_zero)
input_bits = np.iinfo(cast.input_dtype).bits
input_quant_result_tensor_type = \
attributes.QuantResultTensorType(type=TensorType(scalar=input_scalar_type, shape=cast.shape),
quant=Quantization.representable(cast.scale, cast.zero_point, input_bits),
requant_method=RequantMethod.fractional_zero)
calib_attrs = attributes.AwesomeCalibAttrs(
quant=TensorValue(quant_result_tensor_type),
input_quant={InputName("data"): TensorValue(input_quant_result_tensor_type)}
)
return AwesomeNode(name=NodeName(f"dequantize_{name_counter}"),
input_names=[InputName("data")],
input_node_names=[input_name],
ir=_create_sima_ir(operations.DequantizationTransformOp(), attrs,
calib_attrs, backend),
_status=Status.SIMA_QUANTIZED)
[docs]
def create_requantization_node(input_name: NodeName, name_counter: int, input_type: TensorType,
input_quant: Quantization, output_quant: Quantization,
requant_method: RequantMethod,
requant: BaseRequantization[np.ndarray]) \
-> AwesomeNode:
"""
Create requantization node that will be used for converting data from int32 to int16 or int8 type.
:param input_name: Input node name
:param name_counter: Node index
:param input_type: TensorType of the input tensor
:param input_quant: Quantization of input tensor
:param output_quant: Quantization of output tensor
:param requant_method: Requantization method
:param requant: Requantization to perform
:return: AwesomeNode
"""
out_dtype = requant.out_dtype
attrs = attributes.RequantizeAttrs(axis=1, rounding="TOEVEN", compute_dtype='int32',
out_dtype=out_dtype.__name__, input_type=input_type)
quant_attrs = attributes.RequantizeQuantAttrs(attrs=attrs, requant=requant)
quant_result_tensor_type = \
attributes.QuantResultTensorType(type=TensorType(scalar=ScalarType.from_numpy(out_dtype),
shape=input_type.shape),
quant=output_quant, requant_method=RequantMethod.arith_folded)
input_quant_result_tensor_type = \
attributes.QuantResultTensorType(type=input_type, quant=input_quant, requant_method=requant_method)
calib_attrs = attributes.AwesomeCalibAttrs(
quant=TensorValue(quant_result_tensor_type),
input_quant={InputName("data"): TensorValue(input_quant_result_tensor_type)}
)
return AwesomeNode(
name=NodeName(f'requantize_{name_counter}'),
input_names=[InputName('data')],
input_node_names=[input_name],
ir=_create_sima_ir(operations.RequantizeOp(), quant_attrs, calib_attrs, Backend.NONE),
_status=Status.SIMA_QUANTIZED
)
[docs]
def create_cast_node(input_name: NodeName, name_counter: int, shape: Tuple[int, ...],
input_type: ScalarType, output_type: ScalarType) -> AwesomeNode:
"""
Create a node that casts tensors from one scalar type to another.
Casting does not requantize, but only converts the data to a different numeric type.
:param input_name: Input of the new cast node
:param name_counter: Node index
:param shape: Shape of the input and output
:param input_type: Scalar type of the input
:param output_type: Scalar type of the output
:return: AwesomeNode
"""
out_dtype_name = output_type.numpy_type().__name__
in_t = TensorType(input_type, shape)
out_t = TensorType(output_type, shape)
attrs = attributes.CastAttrs(out_dtype=out_dtype_name, input_shape=shape, input_type=input_type)
calib_attrs = attributes.AwesomeCalibAttrs(
quant=TensorValue(attributes.QuantResultTensorType.from_type(out_t)),
input_quant={InputName("data"): TensorValue(attributes.QuantResultTensorType.from_type(in_t))}
)
return AwesomeNode(
name=NodeName(f'cast_{name_counter}'),
input_names=[InputName('data')],
input_node_names=[input_name],
ir=_create_sima_ir(operations.CastOp(), attrs, calib_attrs, Backend.NONE),
_status=Status.SIMA_QUANTIZED
)