float16.py

# MIT License
#
# Copyright (c) Microsoft Corporation, Hugging Face. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


from typing import Optional
import itertools
import numpy as np
import onnx
import packaging.version as pv
import warnings
from onnx import helper, numpy_helper
from onnx import onnx_pb as onnx_proto
import onnxslim.third_party.onnx_graphsurgeon as gs


FLOAT32 = 1
FLOAT16 = 10


def _npfloat16_to_int(np_list):
    """
    Convert numpy float16 to python int.

    :param np_list: numpy float16 list
    :return int_list: python int list
    """
    return [int(bin(_.view("H"))[2:].zfill(16), 2) for _ in np_list]


def convert_np_to_float16(np_array, min_positive_val=1e-7, max_finite_val=1e4):
    """
    Convert float32 numpy array to float16 without changing sign or finiteness.
    Positive values less than min_positive_val are mapped to min_positive_val.
    Positive finite values greater than max_finite_val are mapped to max_finite_val.
    Similar for negative values. NaN, 0, inf, and -inf are unchanged.
    """

    def between(a, b, c):
        return np.logical_and(a < b, b < c)

    positive_values = np_array[np.where(np_array > 0)]
    if positive_values.shape[0] > 0:
        pos_max = positive_values.max()
        pos_min = positive_values.min()

        if pos_max >= max_finite_val:
            warnings.warn(
                "the float32 number {} will be truncated to {}".format(
                    pos_max, max_finite_val
                )
            )

        if pos_min <= min_positive_val:
            warnings.warn(
                "the float32 number {} will be truncated to {}".format(
                    pos_min, min_positive_val
                )
            )

    negative_values = np_array[np.where(np_array < 0)]
    if negative_values.shape[0] > 0:
        neg_max = negative_values.max()
        neg_min = negative_values.min()

        if neg_min <= -max_finite_val:
            warnings.warn(
                "the float32 number {} will be truncated to {}".format(
                    neg_min, -max_finite_val
                )
            )

        if neg_max >= -min_positive_val:
            warnings.warn(
                "the float32 number {} will be truncated to {}".format(
                    neg_max, -min_positive_val
                )
            )

    np_array = np.where(
        between(0, np_array, min_positive_val), min_positive_val, np_array
    )
    np_array = np.where(
        between(-min_positive_val, np_array, 0), -min_positive_val, np_array
    )
    np_array = np.where(
        between(max_finite_val, np_array, float("inf")), max_finite_val, np_array
    )
    np_array = np.where(
        between(float("-inf"), np_array, -max_finite_val), -max_finite_val, np_array
    )
    return np.float16(np_array)


def convert_tensor_float_to_float16(tensor, min_positive_val=1e-7, max_finite_val=1e4):
    """
    Convert tensor float to float16.

    :param tensor: TensorProto object
    :return tensor_float16: converted TensorProto object
    """
    if not isinstance(tensor, onnx_proto.TensorProto):
        raise ValueError(
            "Expected input type is an ONNX TensorProto but got %s" % type(tensor)
        )

    if tensor.data_type == onnx_proto.TensorProto.FLOAT:
        tensor.data_type = onnx_proto.TensorProto.FLOAT16
        # convert float_data (float type) to float16 and write to int32_data
        if tensor.float_data:
            float16_data = convert_np_to_float16(
                np.array(tensor.float_data), min_positive_val, max_finite_val
            )
            int_list = _npfloat16_to_int(float16_data)
            tensor.int32_data[:] = int_list
            tensor.float_data[:] = []
        # convert raw_data (bytes type)
        if tensor.raw_data:
            # convert n.raw_data to float
            float32_list = np.fromstring(tensor.raw_data, dtype="float32")
            # convert float to float16
            float16_list = convert_np_to_float16(
                float32_list, min_positive_val, max_finite_val
            )
            # convert float16 to bytes and write back to raw_data
            tensor.raw_data = float16_list.tostring()
    return tensor


def make_value_info_from_tensor(tensor):
    shape = numpy_helper.to_array(tensor).shape
    return helper.make_tensor_value_info(tensor.name, tensor.data_type, shape)


DEFAULT_OP_BLOCK_LIST = [
    "ArrayFeatureExtractor",
    "Binarizer",
    "CastMap",
    "CategoryMapper",
    "DictVectorizer",
    "FeatureVectorizer",
    "Imputer",
    "LabelEncoder",
    "LinearClassifier",
    "LinearRegressor",
    "Normalizer",
    "OneHotEncoder",
    "RandomUniformLike",
    "SVMClassifier",
    "SVMRegressor",
    "Scaler",
    "TreeEnsembleClassifier",
    "TreeEnsembleRegressor",
    "ZipMap",
    "NonMaxSuppression",
    "TopK",
    "RoiAlign",
    "Resize",
    # 'Range',
    "CumSum",
    "Min",
    "Max",
    "Upsample",
    # NEW:
    "RandomNormalLike",
    # TODO: Ideally, "Cast" nodes should not be here, for the following reasons:
    #  - It breaks the semantics that the default list contains "ops that are not supported for float16 in ONNX Runtime".
    #  - When fp32 casts already exist in the model (e.g., for rotary embeddings), this script will insert redundant casts around it.
    # However, without it, the graphs produced are invalid. Eventually, we will resolve this.
    "Cast",
]


def initial_checking(model, disable_shape_infer):
    func_infer_shape = None
    if not disable_shape_infer and pv.Version(onnx.__version__) >= pv.Version("1.2"):
        try:
            from onnx.shape_inference import infer_shapes

            func_infer_shape = infer_shapes
        finally:
            pass

    if not isinstance(model, onnx_proto.ModelProto):
        raise ValueError(
            "Expected model type is an ONNX ModelProto but got %s" % type(model)
        )

    if func_infer_shape is not None:
        model = func_infer_shape(model)

    is_fp16_ready_flag = check_if_fp16_ready(model.graph)

    return model, func_infer_shape, is_fp16_ready_flag


def convert_float_to_float16(
    model,
    min_positive_val=1e-7,
    max_finite_val=1e4,
    keep_io_types=False,
    disable_shape_infer=False,
    op_block_list=None,
    node_block_list=None,
    check_fp16_ready=True,
):

    # create blocklists
    if op_block_list is None:
        op_block_list = DEFAULT_OP_BLOCK_LIST
    if node_block_list is None:
        node_block_list = []
    op_block_list = set(op_block_list)
    node_block_list = set(node_block_list)

    global_input_name_dict = (
        {}
    )  # key: input name, value: new output name after Cast node
    # basic checking, including shape inference
    model, func_infer_shape, is_fp16_ready_flag = initial_checking(
        model, disable_shape_infer
    )
    if is_fp16_ready_flag and check_fp16_ready:
        raise ValueError(
            "The model is already converted to float16, if convert again, the model might be wrong. \n If you are sure to convert again, please set check_fp16_ready=False."
        )

    graph_stack = [model.graph]

    is_top_level = True
    while graph_stack:
        next_level = []
        for curr_graph in graph_stack:
            process_graph_input(
                curr_graph, is_top_level, keep_io_types, global_input_name_dict
            )
            value_info_block_list = process_tensor_in_node(
                curr_graph,
                op_block_list,
                node_block_list,
                min_positive_val,
                max_finite_val,
            )
            process_value_info(curr_graph, value_info_block_list)
            process_node_in_block_list(
                curr_graph, global_input_name_dict, op_block_list, node_block_list
            )
            process_initializers(
                curr_graph,
                op_block_list,
                node_block_list,
                min_positive_val,
                max_finite_val,
            )
            process_graph_output(curr_graph, is_top_level, keep_io_types)
            sub_graph_list = get_next_level_graph(
                curr_graph, op_block_list, node_block_list
            )
            if len(sub_graph_list) > 0:
                next_level.extend(sub_graph_list)

            if not is_top_level:
                process_node_input_output(curr_graph, global_input_name_dict)
            is_top_level = False  # Going to process sub-graph
        graph_stack = next_level

    remove_unnecessary_cast_node(model.graph)

    # Topologically sort the graph
    # NOTE: We do not perform another round of optimization as the model is already optimized
    graph = gs.import_onnx(model)
    graph.toposort()
    model = gs.export_onnx(graph)

    return model


# Change the input/output of the node to the new output name after Cast node for sub-graph
# Because there have NO value_info start from
def process_node_input_output(
    graph: onnx_proto.GraphProto, global_input_name_dict: dict
):
    for node in graph.node:
        for i, input_name in enumerate(node.input):
            if input_name in global_input_name_dict:
                node.input[i] = global_input_name_dict[input_name]
        for i, output_name in enumerate(node.output):
            if output_name in global_input_name_dict:
                node.output[i] = global_input_name_dict[output_name]


def process_graph_input(
    graph: onnx_proto.GraphProto,
    is_top_level: bool,
    is_io_fp32: bool,
    global_input_name_dict: dict,
):
    # The input dtype is float32, need to cast to fp16
    if is_top_level and is_io_fp32:
        for graph_input in graph.input:  # n_input is ValueInfoProto
            if graph_input.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:
                downstream_nodes = find_downstream_node_by_input_name(
                    graph, graph_input.name
                )
                for d_node in downstream_nodes:
                    # More than one node may consume the model input, so we only create
                    # a single cast node, and then reuse this node when needed.
                    cast_exists = graph_input.name in global_input_name_dict
                    if cast_exists:
                        cast_node_output_name = global_input_name_dict[graph_input.name]
                    else:
                        cast_node_output_name = graph_input.name + "_fp16"
                        add_cast_node(
                            graph,
                            [graph_input.name],
                            [cast_node_output_name],
                            cast_node_output_name,  # Set node name same as output name
                            FLOAT16,
                        )
                        add_new_value_info(
                            graph,
                            graph_input,
                            cast_node_output_name,
                            onnx_proto.TensorProto.FLOAT16,
                        )
                    for i, input_name in enumerate(d_node.input):
                        if input_name == graph_input.name:
                            d_node.input[i] = (
                                cast_node_output_name  # Change the input of the second node
                            )
                            global_input_name_dict[graph_input.name] = (
                                cast_node_output_name
                            )

    # For the sub-graph, don't do cast
    else:  # Change the input dtype to fp16 without any cast
        for graph_input in graph.input:
            if graph_input.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:
                graph_input.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16


def process_graph_output(
    graph: onnx_proto.GraphProto, is_top_level: bool, is_io_fp32: bool
):
    if is_top_level and is_io_fp32:  # the output dtype is float32, need to cast to fp16
        for i, graph_output in enumerate(graph.output):
            if graph_output.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:
                new_producer_name = graph_output.name + "_fp16"
                original_name = graph_output.name  # The correct output name

                # Get the node(s) that produce the model output
                # These will most likely be fp16, but could be fp32 if the previous node is in block_list
                upstream_nodes = find_upstream_node_by_output_name(graph, original_name)
                assert len(upstream_nodes) == 1  # Should be only one node

                producer_node = upstream_nodes[0]

                for i, output_name in enumerate(producer_node.output):
                    if output_name == original_name:
                        producer_node.output[i] = new_producer_name

                cast_node_name = new_producer_name + "_input_cast" + str(i)
                add_cast_node(
                    graph,
                    [new_producer_name],
                    [original_name],
                    cast_node_name,
                    onnx_proto.TensorProto.FLOAT,
                )
                for value_info in graph.value_info:
                    if original_name == value_info.name:
                        value_info.type.tensor_type.elem_type = (
                            onnx_proto.TensorProto.FLOAT
                        )

                # Get the node(s) that consume the model output
                downstream_nodes = find_downstream_node_by_input_name(
                    graph,
                    original_name,
                    include_subgraphs=False,
                )

                # It is possible that the producer node is also input to downstream nodes
                # So, we update the inputs of these downstream nodes
                for d_node in downstream_nodes:
                    for i, input_name in enumerate(d_node.input):
                        if input_name == original_name:
                            d_node.input[i] = new_producer_name

    else:  # change the output dtype to fp16 in tensor
        for graph_output in graph.output:
            if graph_output.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:
                graph_output.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16


def process_node_in_block_list(
    graph: onnx_proto.GraphProto,
    global_input_name_dict: dict,
    op_block_list: list,
    node_block_list: list,
):
    # NB: Important to create a copy of the nodes in the graph to avoid modifying
    # the graph in-place while iterating (causing an infinite loop)
    for node in list(graph.node):
        if (node.op_type in op_block_list) or (node.name in node_block_list):
            insert_cast32_before_node(graph, node, global_input_name_dict)
            insert_cast16_after_node(graph, node, global_input_name_dict)


# Todo: global_input_name_dict still not fill value
def insert_cast32_before_node(
    graph: onnx_proto.GraphProto, node: onnx_proto.NodeProto, global_input_name_dict
):
    for i, input_name in enumerate(node.input):
        for value_info in itertools.chain(graph.value_info, graph.input):
            if input_name == value_info.name:
                if (
                    value_info.type.tensor_type.elem_type
                    != onnx_proto.TensorProto.FLOAT16
                ):
                    break
                cast_output_name = node.name + "_input_cast_" + str(i)
                add_new_value_info(
                    graph, value_info, cast_output_name, onnx_proto.TensorProto.FLOAT
                )
                cast_node_name = node.name + "_input_cast" + str(i)
                add_cast_node(
                    graph,
                    [input_name],
                    [cast_output_name],
                    cast_node_name,
                    onnx_proto.TensorProto.FLOAT,
                )
                node.input[i] = cast_output_name
                break


# Todo: global_input_name_dict still not fill value
def insert_cast16_after_node(
    graph: onnx_proto.GraphProto, node: onnx_proto.NodeProto, global_input_name_dict
):
    for i, output_name in enumerate(node.output):
        for value_info in itertools.chain(graph.value_info, graph.output):
            if output_name == value_info.name:
                if (
                    value_info.type.tensor_type.elem_type
                    != onnx_proto.TensorProto.FLOAT
                ):
                    break
                cast_input_name = node.name + "_output_cast_" + str(i)
                add_new_value_info(
                    graph, value_info, cast_input_name, onnx_proto.TensorProto.FLOAT
                )
                value_info.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16
                cast_node_name = node.name + "_output_cast" + str(i)
                add_cast_node(
                    graph,
                    [cast_input_name],
                    [output_name],
                    cast_node_name,
                    onnx_proto.TensorProto.FLOAT16,
                )
                node.output[i] = cast_input_name
                break


# Process tensor data in attribute of the node
def process_tensor_in_node(
    graph: onnx_proto.GraphProto,
    op_block_list: list,
    node_block_list: list,
    min_positive_val,
    max_finite_val,
):
    value_info_block_list = set()  # This is for later use, not in this step
    for node in graph.node:
        # NOTE: "Cast" operation cannot change its output type because it is strongly typed.
        if (
            (node.op_type in op_block_list)
            or (node.name in node_block_list)
            or (node.op_type == "Cast")
        ):
            # if (node.op_type in op_block_list) or (node.name in node_block_list):
            # Only need to block the output value_info changing
            for output_name in node.output:
                value_info_block_list.add(output_name)
        else:
            for attr in node.attribute:
                # one tensor
                if attr.t.data_type == onnx_proto.TensorProto.FLOAT:
                    attr.t.CopyFrom(
                        convert_tensor_float_to_float16(
                            attr.t, min_positive_val, max_finite_val
                        )
                    )
                # list of tensor
                for t in attr.tensors:
                    if t.data_type == onnx_proto.TensorProto.FLOAT:
                        t.CopyFrom(
                            convert_tensor_float_to_float16(
                                t, min_positive_val, max_finite_val
                            )
                        )
    return value_info_block_list


# Change all the value info type from float32 to float16 if not in block list
def process_value_info(graph: onnx_proto.GraphProto, value_info_block_list: list):
    for value_info in graph.value_info:
        if value_info.name in value_info_block_list:
            continue
        else:
            if value_info.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT:
                value_info.type.tensor_type.elem_type = onnx_proto.TensorProto.FLOAT16


# Initializer is 'edge' type, so doesn't have value_info
def process_initializers(
    graph: onnx_proto.GraphProto,
    op_block_list,
    node_block_list,
    min_positive_val,
    max_finite_val,
):
    # Find the input of the block node, don't need to change this kind of initializer
    initializer_block_list = set()
    for node in graph.node:
        if (node.op_type in op_block_list) or (node.name in node_block_list):
            for (
                input_name
            ) in (
                node.input
            ):  # some is initializer, some is value_info, can't distinguish but doesn't matter
                initializer_block_list.add(input_name)
    # Process initializers
    for initializer in graph.initializer:
        if initializer.name not in initializer_block_list:
            if initializer.data_type == onnx_proto.TensorProto.FLOAT:
                convert_tensor_float_to_float16(
                    initializer, min_positive_val, max_finite_val
                )


def get_next_level_graph(
    graph: onnx_proto.GraphProto, op_block_list: list, node_block_list: list
):
    sub_graph_list = []
    for node in graph.node:
        if node.op_type in op_block_list or node.name in node_block_list:
            continue
        for attr in node.attribute:
            # Check if sub-graph exist
            if len(attr.g.node) > 0:  # single sub-graph
                sub_graph_list.append(attr.g)
            for g in attr.graphs:
                if len(g.node) > 0:  # multiple sub-graphs
                    sub_graph_list.append(g)
    return sub_graph_list


def add_cast_node(
    graph: onnx_proto.GraphProto,
    inputs: list,
    outputs: list,
    node_name: str,
    to_type: int,
):
    new_node = [helper.make_node("Cast", inputs, outputs, to=to_type, name=node_name)]
    graph.node.extend(new_node)


def add_new_value_info(
    graph: onnx_proto.GraphProto,
    exist_value_info: onnx_proto.ValueInfoProto,
    name: str,
    dtype: int,
):
    new_value_info = graph.value_info.add()
    new_value_info.CopyFrom(exist_value_info)
    new_value_info.name = name
    new_value_info.type.tensor_type.elem_type = dtype


# Find the node that has the specified output name
def find_upstream_node_by_output_name(graph: onnx_proto.GraphProto, output_name: str):
    nodes = []
    for node in graph.node:
        if output_name in node.output:
            nodes.append(node)
    assert len(nodes) <= 1  # Suppose there is less than one node found
    return nodes


# Find the node that has the specified input name, including in subgraphs
def find_downstream_node_by_input_name(
    graph: onnx_proto.GraphProto, input_name: str, include_subgraphs=True
):
    nodes = []

    # Check nodes in current graph
    for node in graph.node:
        if input_name in node.input:
            nodes.append(node)

        if not include_subgraphs:
            continue

        # Recursively check subgraphs in node attributes
        for attr in node.attribute:
            if attr.type == onnx_proto.AttributeProto.GRAPH:
                # Single subgraph
                if len(attr.g.node) > 0:
                    nodes.extend(find_downstream_node_by_input_name(attr.g, input_name))

            # Multiple subgraphs
            if attr.type == onnx_proto.AttributeProto.GRAPHS:
                for g in attr.graphs:
                    if len(g.node) > 0:
                        nodes.extend(find_downstream_node_by_input_name(g, input_name))

    return nodes


# Remove identity node
def remove_identity_node_from_model(model: onnx_proto.ModelProto):
    remove_identity_node_from_graph(model.graph)
    try:
        from onnx.shape_inference import infer_shapes

        func_infer_shape = infer_shapes
        model = func_infer_shape(model)
        return model
    finally:
        pass


# Remove identity node
def remove_identity_node_from_graph(graph: onnx_proto.GraphProto):
    for curr_node in graph.node:
        if curr_node.op_type == "Identity":
            for input_name in curr_node.input:
                upstream_nodes = find_upstream_node_by_output_name(graph, input_name)
                for u_node in upstream_nodes:
                    if u_node is not None:
                        u_node.output[0] = curr_node.output[0]
                        graph.node.remove(curr_node)


def convert_float_to_float16_model_path(
    model_path, min_positive_val=1e-7, max_finite_val=1e4, keep_io_types=False
):
    """
    Convert tensor float type in the ONNX Model to tensor float16.
    *It is to fix an issue that infer_shapes func cannot be used to infer >2GB models.
    *But this function can be applied to all model sizes.
    :param model_path: ONNX Model path
    :return: converted ONNX ModelProto object
    Examples
    ::
        #Convert to ONNX ModelProto object and save model binary file:
        from onnxmltools.utils.float16_converter import convert_float_to_float16_model_path
        new_onnx_model = convert_float_to_float16_model_path('model.onnx')
        onnx.save(new_onnx_model, 'new_model.onnx')
    """

    disable_shape_infer = False
    if pv.Version(onnx.__version__) >= pv.Version("1.8"):
        try:
            # infer_shapes_path can be applied to all model sizes
            from onnx.shape_inference import infer_shapes_path
            import tempfile
            import os

            # shape_infer_model_path should be in the same folder of model_path
            with tempfile.NamedTemporaryFile(
                dir=os.path.dirname(model_path)
            ) as tmpfile:
                shape_infer_model_path = tmpfile.name
                infer_shapes_path(model_path, shape_infer_model_path)
                model = onnx.load(shape_infer_model_path)
                disable_shape_infer = True
        finally:
            pass
    if not disable_shape_infer:
        model = onnx.load(model_path)
    return convert_float_to_float16(
        model, min_positive_val, max_finite_val, keep_io_types, disable_shape_infer
    )


def remove_unnecessary_cast_node(graph_proto: onnx_proto.GraphProto):
    # 1. find all cast nodes in the graph
    cast_node_list = []
    input_name_to_cast_node_dict = {}
    output_name_to_cast_node_dict = {}
    # using name as key to point to a node. because node object cannot be key
    name_to_node_dict = {}
    for node in graph_proto.node:
        if node.op_type == "Cast":
            # if node.name not in ["graph_input_cast0", "graph_output_cast0"]:
            cast_node_list.append(node)

            name_to_node_dict[node.name] = node
            for input_name in node.input:
                input_name_to_cast_node_dict[input_name] = node
            for output_name in node.output:
                output_name_to_cast_node_dict[output_name] = node

    # 2. find upstream and downstream node of the cast node
    cast_node_upstream_dict = {}  # mapping cast node(name) to its upstream node
    cast_node_downstream_dict = {}  # mapping cast node(name) to its downstream node
    for current_node in graph_proto.node:
        # find the downstream node(s)
        for input_name in current_node.input:
            if input_name in output_name_to_cast_node_dict:
                # found the downstream node of the cast node, might be multiple
                cast_node = output_name_to_cast_node_dict[input_name]
                if cast_node.name not in cast_node_downstream_dict:
                    cast_node_downstream_dict[cast_node.name] = current_node
                else:  # already exists one downstream node, make it a list
                    existing_downstream_nodes = cast_node_downstream_dict[
                        cast_node.name
                    ]
                    if isinstance(existing_downstream_nodes, list):
                        existing_downstream_nodes.append(current_node)
                    else:  # make a list
                        existing_downstream_nodes = [
                            existing_downstream_nodes,
                            current_node,
                        ]
                        cast_node_downstream_dict[cast_node.name] = (
                            existing_downstream_nodes
                        )
        # find the upstream node
        for output_name in current_node.output:
            if output_name in input_name_to_cast_node_dict:
                # found the upstream node of the cast node, should be unique
                cast_node = input_name_to_cast_node_dict[output_name]
                cast_node_upstream_dict[cast_node.name] = current_node

    # 3. remove the cast node which upstream is 'Constant'
    for cast_node_name, upstream_node in cast_node_upstream_dict.items():
        cast_node = name_to_node_dict[cast_node_name]
        if upstream_node.op_type == "Constant":
            cast_node_list.remove(cast_node)

    # 4. find (cast_to_fp16, cast_to_fp32) pairs where --fp32--> cast_to_fp16 --fp16--> cast_to_fp32.
    remove_candidate = []

    name_to_value_info = {
        value_info.name: value_info
        for value_info in itertools.chain(graph_proto.value_info, graph_proto.input)
    }

    def get_type(name: str) -> Optional[int]:
        if name in name_to_value_info:
            return name_to_value_info[name].type
        else:
            # `name` has no value info.
            return None

    for cast_node_name, downstream_node in cast_node_downstream_dict.items():
        cast_node = name_to_node_dict[cast_node_name]
        if len(cast_node.input) != 1:
            raise RuntimeError(
                f"Cast node {cast_node_name} should have only one input, but has {len(cast_node.input)}."
            )

        input_type = get_type(cast_node.input[0])
        if input_type != onnx_proto.TensorProto.FLOAT:
            continue
        if isinstance(downstream_node, list):
            for dn in downstream_node:
                if (
                    dn.op_type == "Cast"
                    and dn.attribute[0].i == 32
                    and cast_node.attribute[0].i == 16
                    and dn in cast_node_list
                    and cast_node in cast_node_list
                ):
                    remove_candidate.append((cast_node, dn))
        else:
            if (
                downstream_node.op_type == "Cast"
                and cast_node.attribute[0].i == FLOAT16
                and downstream_node.attribute[0].i == FLOAT32
                and downstream_node in cast_node_list
                and cast_node in cast_node_list
            ):
                remove_candidate.append((cast_node, downstream_node))

    # 5. change "upstream --fp32--> cast_to_fp16 --fp16--> cast_to_fp32 --fp32--> downstream" to
    # "upstream --fp32--> downstream".
    for cast_node_pair in remove_candidate:
        first_cast_node = cast_node_pair[0]
        second_cast_node = cast_node_pair[1]
        upstream_node = cast_node_upstream_dict.get(first_cast_node.name)
        downstream_node = cast_node_downstream_dict.get(second_cast_node.name)
        if upstream_node is None and downstream_node is not None:
            # The upstream_node should be graph input
            out = first_cast_node.input[0]
            for i, input_name in enumerate(downstream_node.input):
                for output_name in second_cast_node.output:
                    if input_name == output_name:
                        # change the input as the upstream node's output
                        downstream_node.input[i] = out
        elif upstream_node is not None and downstream_node is None:
            raise ValueError(
                "The downstream node of the second cast node should be graph output"
            )
        else:
            # find the upstream node's output to first_cast_node
            out = None
            for output_name in upstream_node.output:
                if output_name == first_cast_node.input[0]:
                    out = output_name
                    break
            # find the downstream node's input as second_cast_node's output
            for i, input_name in enumerate(downstream_node.input):
                for output_name in second_cast_node.output:
                    if input_name == output_name:
                        # change the input as the upstream node's output
                        downstream_node.input[i] = out

    # 6. remove the cast node pair
    for cast_node_pair in remove_candidate:
        graph_proto.node.remove(cast_node_pair[0])
        graph_proto.node.remove(cast_node_pair[1])


# Check if the model is already converted to float16
def check_if_fp16_ready(graph_proto):
    # Check graph input and ouput
    is_value_info_fp16 = False
    for value_info in itertools.chain(
        graph_proto.output, graph_proto.input, graph_proto.value_info
    ):
        if value_info.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT16:
            is_value_info_fp16 = True
            break

    # Check initializer
    is_initializer_fp16 = False
    for initializer in graph_proto.initializer:
        if initializer.data_type == onnx_proto.TensorProto.FLOAT16:
            is_initializer_fp16 = True
            break

    # Check cast node
    has_cast_node_fp16 = False
    for node in graph_proto.node:
        if node.op_type == "Cast" and node.attribute[0].i == FLOAT16:
            has_cast_node_fp16 = True
            break

    # Any of above flags is True, return True
    if is_value_info_fp16 or is_initializer_fp16 or has_cast_node_fp16:
        return True  # already converted to float16
    else:
        return False  # not converted to float16 yet