Source code for afe.ir.node

#########################################################
# Copyright (C) 2020 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
#########################################################
from collections import OrderedDict
from typing import List, Union, Dict, Any, Optional
from dataclasses import dataclass

from afe.backends import BackendIR, Backend
from afe.ir import attributes as afe_attrs
from afe.ir import operations as afe_op
from afe.ir.defines import Status, NodeName, InputName, data_value_elements, LayerStats
from afe.ir.operations import PlaceholderOp
from afe.ir.sima_ir import SiMaIR
from afe.ir.tensor_type import NodeType, ScalarType


@dataclass


class AwesomeNode:
    """
    An awesome node is responsible for storing attributes and operations used by the AwesomeNet when
    executing the network

    :param name: Name of the node
    :param input_names: Parameter names of this AwesomeNode.  These represent the inputs expected by
       the code inside this node.
    :param input_node_names: Argument names of this AwesomeNode.  These represent the inputs supplied
       by the net that contains this node.  Data for each argument gets forwarded to the corresponding
       parameter.
    :param ir: Union[SiMaIR, "afe.ir.net.AwesomeNet", "afe.backends.BackendIR"]. The IR can be:
        * SiMaIR - A frontend IR that contains single operation or ExternalOp
        * AwesomeNet - A frontend IR that contains a sub-graph that can be supported by SiMa's backends
        * BackendIR - A backend IR that can be compiled, or already compiled by certain backend

    :param _status: Status. Node status. Default is Status.RELAY
    :param _layer_stats: Layer statistics. For each MLA node, quantization error is calculated,
                         that information is than forwarded to .sima.json file, and it can be viewed in Netron.
    """


    name: NodeName



    input_names: List[InputName]



    input_node_names: List[NodeName]



    ir: Union[SiMaIR, "afe.ir.net.AwesomeNet", "afe.backends.BackendIR"]


    _status: Status = Status.RELAY
    _layer_stats: Optional[LayerStats] = None

    @property


    def status(self) -> Status:
        return self._status


    @status.setter
    def status(self, status: Status):
        self._status = status



    def get_type(self) -> NodeType:
        import afe.ir.net
        import afe.backends

        if isinstance(self.ir, (SiMaIR, afe.backends.BackendIR)):
            return self.ir.get_type()
        elif isinstance(self.ir, afe.ir.net.AwesomeNet):
            return _make_net_type(self.input_names, self.ir)
        else:
            raise TypeError("Unrecognized type of data in AwesomeNode")




    def set_batch_size(self, batch_size: int):
        """
        Modifies AwesomeNode'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.ir.set_batch_size(batch_size)




    def update_layer_stats(self, layer_stats: LayerStats):
        self._layer_stats = layer_stats


    @property


    def layer_stats(self) -> LayerStats:
        return self._layer_stats




def _make_net_type(input_names: List[NodeName], net: "afe.ir.net.AwesomeNet") -> NodeType:
    """
    Make the type of an AwesomeNode that contains an AwesomeNet.
    The parameter names are taken from the AwesomeNode.  Their associated types and
    the return type are taken from the AwesomeNet.

    :param input_names: The AwesomeNode's parameters.  These determine the parameters
       in the node type.
    :param net: The AwesomeNet that implements the node.
    :return: The type of the node.
    """
    assert len(input_names) == len(net.input_node_names)
    parameter_types = {name: net.nodes[input_name].get_type().output
                       for name, input_name in zip(input_names, net.input_node_names)}
    return_type = net.nodes[net.output_node_name].get_type().output

    return NodeType(parameter_types, return_type)




def node_is_sima_ir(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a SiMa IR

    :param node: AwesomeNode
    """
    return isinstance(node.ir, SiMaIR)





def node_is_backend_ir(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a SiMa IR

    :param node: AwesomeNode
    """
    return isinstance(node.ir, BackendIR)





def node_is_constant(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a Constant SiMaIR

    :param node: AwesomeNode
    """
    return node_is_sima_ir(node) and isinstance(node.ir.operation, afe_op.ConstantOp)





def node_is_placeholder(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a Placeholder SiMaIR

    :param node: AwesomeNode
    """
    return node_is_sima_ir(node) and isinstance(node.ir.operation, PlaceholderOp)





def node_is_awesomenet(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a AwesomeNet

    :param node: AwesomeNode
    """
    return node.ir.__class__.__name__ == "AwesomeNet"





def node_is_external(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a External Graph

    :param node: AwesomeNode
    """
    return node_is_sima_ir(node) and isinstance(node.ir.attrs, afe_attrs.ExternalAttrs)





def node_is_ev(node: AwesomeNode) -> bool:
    """
    Return True if the AwesomeNode is an EV node.

    :param node: AwesomeNode
    """
    return node_is_sima_ir(node) and node.ir.backend == Backend.EV





def node_is_relu(node: AwesomeNode) -> bool:
    return (node_is_sima_ir(node)
            and isinstance(node.ir.operation, (afe_op.ReluOp, afe_op.PReluOp, afe_op.LeakyReluCompositeOp)))





def node_is_mean(node: AwesomeNode) -> bool:
    return node_is_sima_ir(node) and isinstance(node.ir.operation, afe_op.MeanOp)





def node_is_upsampling(node: AwesomeNode) -> bool:
    return node_is_sima_ir(node) and isinstance(node.ir.operation, (afe_op.UpsamplingOp, afe_op.ImageResize2DOp))





def node_is_subgraph(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a sub-graph.

    Parameters
    ----------
    :param node: AwesomeNode

    Return
    ------
    :return: bool
    """
    return node_is_awesomenet(node) or node_is_external(node)





def node_is_tuple(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a Tuple.

    Parameters
    ----------
    :param node: AwesomeNode

    Return
    ------
    :return: bool
    """
    return node_is_sima_ir(node) and isinstance(node.ir.attrs, afe_attrs.TupleAttrs)





def node_is_tuple_get_item(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains a TupleGetItem.

    Parameters
    ----------
    :param node: AwesomeNode

    Return
    ------
    :return: bool
    """
    return node_is_sima_ir(node) and isinstance(node.ir.attrs, afe_attrs.TupleGetItemAttrs)





def node_is_unpack(node: AwesomeNode) -> bool:
    """
    Return True is the AwesomeNode contains an Unpack transform.

    Parameters
    ----------
    :param node: AwesomeNode

    Return
    ------
    :return: bool
    """
    return node_is_sima_ir(node) and isinstance(node.ir.attrs, afe_attrs.UnpackTransformAttrs)





def node_is_fp32_node(node: AwesomeNode) -> bool:
    """
    Return True if AwesomeNode's output type is floating point.

    :param node: AwesomeNode to be analyzed.
    :return: bool. True if AwesomeNode's output type is fp32, otherwise False.
    """
    output_type = node.get_type().output
    return all([el.scalar == ScalarType.float32 for el in data_value_elements(output_type)])





def node_is_pool_node(node: AwesomeNode) -> bool:
    """
    Return True if the AwesomeNode contains Pool operator.

    :param node: AwesomeNode to be analyzed.
    :return: bool. True if the AwesomeNode's operator is Pooling operator (MaxPool, AvgPool).
    """
    return node_is_sima_ir(node) and isinstance(node.ir.operation, (afe_op.AvgPool2DOp, afe_op.MaxPool2DOp))





def node_is_requantization_node(node: AwesomeNode) -> bool:
    """
    Return True if the AwesomeNode contains Requantize operator.

    :param node: AwesomeNode to be analyzed.
    :return: bool. True if the AwesomeNode's operator is Requantize operator.
    """
    return node_is_sima_ir(node) and isinstance(node.ir.operation, afe_op.RequantizeOp)





def node_uses_observer(node: AwesomeNode) -> bool:
    """
    Return True if node uses calibration information during quantization.

    :param node: AwesomeNode to be analyzed.
    :return: bool. True if AwesomeNode uses calibration information during quantization.
        Otherwise, quantization information is obtained using quantization from its input
        node(s).
    """
    # Currently, nodes that do not use calibration information and thus do not contain
    # NodeObservers are:
    #   - nodes that bypass input to output
    #   - external nodes
    #   - pool nodes
    # TODO: Extend this list?
    return not (
            node.ir.should_bypass_input_to_output() or
            node_is_external(node) or
            node_is_pool_node(node) or
            node_is_mean(node) or
            node_is_relu(node) or
            node_is_upsampling(node)
    )





def get_node_inputs(node: AwesomeNode, node_outputs: Dict[NodeName, Any]) -> Dict[NodeName, Any]:
    """
    Generates a dictionary using the keys from a node's input_names and the values corresponding
    to the data of nodes whose names are in the node's input_node_names.

    :param node: The AwesomeNode we are gathering inputs for.
    :param node_outputs: Dictionary containing outputs from other nodes from across the network.
    :return: A dictionary containing node inputs.
    """
    input_dict = OrderedDict()
    for k, v in zip(node.input_names, node.input_node_names):
        input_dict[k] = node_outputs[v]
    return input_dict