#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2021 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import tensorflow as tf

from tensorflow.python.framework import tensor_util
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.framework import node_def_pb2
from tensorflow.python.framework import dtypes

from ..graph_base import GraphRewriterBase
from neural_compressor.adaptor.tf_utils.graph_util import GraphAnalyzer
from neural_compressor.adaptor.tf_utils.graph_util import GraphRewriterHelper as Helper
from neural_compressor.adaptor.tf_utils.util import version1_gt_version2, version1_lt_version2


class FuseMatMulRequantizeDequantizeTransformer(GraphRewriterBase):
    """Fuse QuantizedMatMul + Requantize + Dequantize into QuantizedMatMulWithBiasAndDequantize.
    """
    def __init__(self, model, device='cpu'):
        super().__init__(model)
        self.device = device
        self.graph_analyzer = GraphAnalyzer()
        self.graph_analyzer.graph = self.model

        self.graph_info = self.graph_analyzer.parse_graph()

        self.eps = 1e-5

    def do_transformation(self):
        fuse_pattern = []
        if tf.version.VERSION in ("1.15.0-up2", "1.15.0-up3") or version1_gt_version2(tf.version.VERSION, '2.1.0'):
            fuse_pattern = [["QuantizedMatMulWithBias"], ['Requantize'], ['Dequantize'], ('Softmax',)]
        float32_type = dtypes.float32.as_datatype_enum
        qint32_type = dtypes.qint32.as_datatype_enum
        target_nodes = self.graph_analyzer.query_fusion_pattern_nodes(fuse_pattern)
        for i in target_nodes:
            # TODO Remove below checker once the TF's limitation removed.
            if len(i) == 5 and version1_lt_version2(tf.__version__, '2.6.0'):
                continue

            quantized_node_name = i[0]
            quantized_node = self.graph_info[quantized_node_name].node
            requantize_node_name = i[1]
            requantize_node = self.graph_info[requantize_node_name].node
            requested_output_min_name = requantize_node.input[3]
            requested_output_max_name = requantize_node.input[4]
            deq_node_name = i[2]

            quantized_node_op = i[-1][0]

            new_node = node_def_pb2.NodeDef()

            new_node.op = quantized_node_op + "AndDequantize"
            new_node.name = requantize_node_name
            for _, value in enumerate(quantized_node.input):
                new_node.input.append(value)

            new_node.input.append(requested_output_min_name)
            new_node.input.append(requested_output_max_name)
            if 'T1' in quantized_node.attr:
                new_node.attr["T1"].CopyFrom(quantized_node.attr['T1'])
            if 'T2' in quantized_node.attr:
                new_node.attr["T2"].CopyFrom(quantized_node.attr['T2'])

            top_node_name = Helper.node_name_from_input(quantized_node.input[0])
            max_filter_node = self.graph_info[new_node.input[6]].node
            min_filter_node = self.graph_info[new_node.input[5]].node
            last_node = self.graph_info[new_node.input[0]].node

            weight_node = self.graph_info[Helper.node_name_from_input(new_node.input[1])].node
            bias_node = self.graph_info[Helper.node_name_from_input(new_node.input[2])].node
            max_input_node = self.graph_info[last_node.input[-1]].node
            min_input_node = self.graph_info[last_node.input[-2]].node

            if max_input_node.op == 'Enter': # pragma: no cover
                min_input_parent_name = Helper.node_name_from_input(min_input_node.input[0])
                max_input_parent_name = Helper.node_name_from_input(max_input_node.input[0])
                min_input_parent_node = self.graph_info[min_input_parent_name].node
                max_input_parent_node = self.graph_info[max_input_parent_name].node
                if min_input_parent_node.op != 'Const' or max_input_parent_node.op != 'Const':
                    continue
                min_input_node = min_input_parent_node
                max_input_node = max_input_parent_node
            if max_filter_node.op == 'Enter': # pragma: no cover
                min_filter_parent_name = Helper.node_name_from_input(min_filter_node.input[0])
                max_filter_parent_name = Helper.node_name_from_input(max_filter_node.input[0])
                min_filter_parent_node = self.graph_info[min_filter_parent_name].node
                max_filter_parent_node = self.graph_info[max_filter_parent_name].node
                if min_filter_parent_node.op != 'Const' or max_filter_parent_node.op != 'Const':
                    continue
                min_filter_node = min_filter_parent_node
                max_filter_node = max_filter_parent_node
            if weight_node.op == 'Enter': # pragma: no cover
                weight_parent_name = Helper.node_name_from_input(weight_node.input[0])
                weight_parent_node = self.graph_info[weight_parent_name].node
                if weight_parent_node.op != 'Const':
                    continue
                weight_node = weight_parent_node
            bias_enter_node = None
            if bias_node.op == 'Enter': # pragma: no cover
                bias_enter_node = bias_node
                bias_parent_name = Helper.node_name_from_input(bias_node.input[0])
                bias_parent_node = self.graph_info[bias_parent_name].node
                if bias_parent_node.op != 'Const':
                    continue
                bias_node = bias_parent_node

            if max_filter_node.op == 'Const':
                min_input_value = (min_input_node.attr['value'].tensor.float_val)[0]
                max_input_value = (max_input_node.attr['value'].tensor.float_val)[0]

                max_filter_value = (max_filter_node.attr['value'].tensor.float_val)[0]
                min_filter_value = (min_filter_node.attr['value'].tensor.float_val)[0]

                weights_tensor = tensor_util.MakeNdarray(
                        weight_node.attr['value'].tensor)
                bias_tensor = tensor_util.MakeNdarray(
                        bias_node.attr['value'].tensor)
                is_min_first = bool(quantized_node.attr['input_quant_mode'].s == b'MIN_FIRST')
                input_range = max_input_value - min_input_value if is_min_first else max(
                        abs(max_input_value), abs(min_input_value))

                if  -self.eps <= input_range <= self.eps:
                    input_range += self.eps

                if -self.eps <= max_input_value - min_input_value <= self.eps:
                    max_input_value += self.eps
                int32_bias = Helper.generate_int32_bias_for_matmul(bias_tensor, weights_tensor,
                                                        input_range, max_input_value,
                                                        min_input_value,
                                                        max_filter_value, min_filter_value)

                bias_node.attr['dtype'].CopyFrom(
                    attr_value_pb2.AttrValue(
                        type=float32_type if self.device == 'gpu' else qint32_type))
                bias_node.attr['value'].CopyFrom(
                    attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
                        bias_tensor if self.device == 'gpu' else int32_bias, dtypes.
                        float32 if self.device == 'gpu' else dtypes.int32, bias_tensor.shape)))

                bias_node.attr['value'].tensor.dtype = float32_type \
                                        if self.device == 'gpu' else qint32_type
                new_node.attr["Tbias"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type \
                                                if self.device == 'gpu' else qint32_type))
                if bias_enter_node:
                    bias_enter_node.attr["T"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type \
                                                if self.device == 'gpu' else qint32_type))
            else:
                new_node.attr["Tbias"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))
            new_node.attr["Toutput"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))

            self.graph_analyzer.remove_node(requantize_node_name)

            if self.graph_info[deq_node_name].outputs:
                self.graph_analyzer.replace_single_node(
                    new_node, [top_node_name], quantized_node_name,
                    self.graph_info[deq_node_name].outputs, deq_node_name)
                self.graph_analyzer.remove_node(deq_node_name)
            else:
                self.graph_analyzer.remove_node(deq_node_name)

                new_node.name = deq_node_name
                self.graph_analyzer.replace_single_node(
                    new_node, [top_node_name], quantized_node_name, [], deq_node_name)

            self.graph_analyzer.remove_node(quantized_node_name)

        return self.graph_analyzer.dump_graph()

class FuseMatMulRequantizeTransformer(GraphRewriterBase):
    """Fuse Quantized MatMul Op with the successor Requantize Op.
    """
    def __init__(self, model, device='cpu'):
        super().__init__(model)
        self.device = device
        self.graph_analyzer = GraphAnalyzer()
        self.graph_analyzer.graph = self.model
        self.eps = 1e-05
        self.graph_info = self.graph_analyzer.parse_graph()

    def do_transformation(self):
        """Fuse the quantized op with the following requantize op.
        Returns:
            [graphdef]: the optimized graphdef object
        """
        uint8_type = dtypes.quint8.as_datatype_enum
        float32_type = dtypes.float32.as_datatype_enum
        qint32_type = dtypes.qint32.as_datatype_enum

        while True:
            target_nodes = self.graph_analyzer.query_fusion_pattern_nodes(
                [["QuantizedMatMulWithBiasAndRelu"], ['Requantize']])
            if len(target_nodes) == 0:
                break

            i = target_nodes[0]
            quantized_node_name = i[0]
            quantized_node = self.graph_info[quantized_node_name].node
            requantize_node_name = i[1]
            requantize_node = self.graph_info[requantize_node_name].node
            requested_output_min_name = requantize_node.input[3]
            requested_output_max_name = requantize_node.input[4]

            quantized_node_op = i[-1][0]

            new_node = node_def_pb2.NodeDef()

            new_node.op = quantized_node_op + "AndRequantize"
            new_node.name = requantize_node_name
            for _, value in enumerate(quantized_node.input):
                new_node.input.append(value)
            new_node.input.append(requested_output_min_name)
            new_node.input.append(requested_output_max_name)
            if 'T1' in quantized_node.attr:
                new_node.attr["T1"].CopyFrom(quantized_node.attr['T1'])
            if 'T2' in quantized_node.attr:
                new_node.attr["T2"].CopyFrom(quantized_node.attr['T2'])

            parent_node_name = Helper.node_name_from_input(quantized_node.input[0])
            max_filter_node = self.graph_info[new_node.input[6]].node
            min_filter_node = self.graph_info[new_node.input[5]].node
            last_node = self.graph_info[new_node.input[0]].node

            is_min_first = bool(quantized_node.attr['input_quant_mode'].s == b'MIN_FIRST')

            weight_node = self.graph_info[new_node.input[1]].node
            bias_node = self.graph_info[new_node.input[2]].node
            max_input_node = self.graph_info[last_node.input[-1]].node
            min_input_node = self.graph_info[last_node.input[-2]].node

            if max_input_node.op == 'Enter': # pragma: no cover
                min_input_parent_name = Helper.node_name_from_input(min_input_node.input[0])
                max_input_parent_name = Helper.node_name_from_input(max_input_node.input[0])
                min_input_parent_node = self.graph_info[min_input_parent_name].node
                max_input_parent_node = self.graph_info[max_input_parent_name].node
                if min_input_parent_node.op != 'Const' or max_input_parent_node.op != 'Const':
                    continue
                min_input_node = min_input_parent_node
                max_input_node = max_input_parent_node
            if max_filter_node.op == 'Enter': # pragma: no cover
                min_filter_parent_name = Helper.node_name_from_input(min_filter_node.input[0])
                max_filter_parent_name = Helper.node_name_from_input(max_filter_node.input[0])
                min_filter_parent_node = self.graph_info[min_filter_parent_name].node
                max_filter_parent_node = self.graph_info[max_filter_parent_name].node
                if min_filter_parent_node.op != 'Const' or max_filter_parent_node.op != 'Const':
                    continue
                min_filter_node = min_filter_parent_node
                max_filter_node = max_filter_parent_node
            if weight_node.op == 'Enter': # pragma: no cover
                weight_parent_name = Helper.node_name_from_input(weight_node.input[0])
                weight_parent_node = self.graph_info[weight_parent_name].node
                if weight_parent_node.op != 'Const':
                    continue
                weight_node = weight_parent_node
            bias_enter_node = None
            if bias_node.op == 'Enter': # pragma: no cover
                bias_enter_node = bias_node
                bias_parent_name = Helper.node_name_from_input(bias_node.input[0])
                bias_parent_node = self.graph_info[bias_parent_name].node
                if bias_parent_node.op != 'Const':
                    continue
                bias_node = bias_parent_node

            if last_node.op.find('Requantize') != -1 or last_node.op.find('QuantizeV2') != -1:
                min_input_value = (min_input_node.attr['value'].tensor.float_val)[0]
                max_input_value = (max_input_node.attr['value'].tensor.float_val)[0]

                max_filter_value = (max_filter_node.attr['value'].tensor.float_val)[0]
                min_filter_value = (min_filter_node.attr['value'].tensor.float_val)[0]

                weights_tensor = tensor_util.MakeNdarray(
                        weight_node.attr['value'].tensor)
                bias_tensor = tensor_util.MakeNdarray(
                        bias_node.attr['value'].tensor)

                input_range = max_input_value - min_input_value if is_min_first else max(
                        abs(max_input_value), abs(min_input_value))

                if -self.eps <= input_range <= self.eps:
                    input_range += self.eps

                if -self.eps <= max_input_value - min_input_value <= self.eps:
                    max_input_value += self.eps
                int32_bias = Helper.generate_int32_bias_for_matmul(bias_tensor, weights_tensor,
                                                                   input_range,
                                                                   max_input_value,
                                                                   min_input_value,
                                                                   max_filter_value,
                                                                   min_filter_value)
                bias_node.attr['dtype'].CopyFrom(
                    attr_value_pb2.AttrValue(
                        type=float32_type if self.device == 'gpu' else qint32_type))
                bias_node.attr['value'].CopyFrom(
                    attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
                        bias_tensor if self.device == 'gpu' else int32_bias, dtypes.
                        float32 if self.device == 'gpu' else dtypes.int32, bias_tensor.shape)))

                bias_node.attr['value'].tensor.dtype = float32_type \
                                        if self.device == 'gpu' else qint32_type
                new_node.attr["Tbias"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type \
                                                if self.device == 'gpu' else qint32_type))

                new_node.attr["Toutput"].CopyFrom(
                        attr_value_pb2.AttrValue(type=uint8_type))
                #TODO enabled below commit once the graph refactor pre_optimize commmitted.
                if quantized_node_op.find('Relu') == -1:
                    deq_node_name = self.graph_info[requantize_node_name].outputs[0]
                    deq_node = self.graph_info[deq_node_name].node
                    deq_node.attr['T'].CopyFrom(attr_value_pb2.AttrValue(type=uint8_type))
                if bias_enter_node:
                    bias_enter_node.attr["T"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type \
                                                if self.device == 'gpu' else qint32_type))
            else:
                new_node.attr["Tbias"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))

            self.graph_analyzer.replace_single_node(
                new_node, [parent_node_name], quantized_node_name,
                [self.graph_info[requantize_node_name].outputs[0]], requantize_node_name)
            self.graph_analyzer.remove_node(quantized_node_name)

        return self.graph_analyzer.dump_graph()

class FuseMatMulRequantizeDequantizeNewAPITransformer(GraphRewriterBase): # pragma: no cover
    """Fuse _QuantizedMatMul + Requantize + Dequantize into _QuantizedMatMul.
    """
    def __init__(self, model, device='cpu'):
        super().__init__(model)
        self.device = device
        self.graph_analyzer = GraphAnalyzer()
        self.graph_analyzer.graph = self.model

        self.graph_info = self.graph_analyzer.parse_graph()

        self.eps = 1e-5

    def do_transformation(self):
        fuse_pattern = [["_QuantizedMatMul"], ['Requantize'], ['Dequantize'], ('Softmax',)]

        uint8_type = dtypes.quint8.as_datatype_enum
        int8_type = dtypes.qint8.as_datatype_enum
        float32_type = dtypes.float32.as_datatype_enum
        qint32_type = dtypes.qint32.as_datatype_enum
        target_nodes = self.graph_analyzer.query_fusion_pattern_nodes(fuse_pattern)
        for i in target_nodes:
            quantized_node_name = i[0]
            quantized_node = self.graph_info[quantized_node_name].node
            requantize_node_name = i[1]
            requantize_node = self.graph_info[requantize_node_name].node
            requested_output_min_name = requantize_node.input[3]
            requested_output_max_name = requantize_node.input[4]
            deq_node_name = i[2]

            quantized_node_op = i[-1][0]
            
            # "BiasAdd" + "Add" only supports "Dequantize"
            attr_fused_ops = ''.join(x for x in quantized_node.attr["fused_ops"].SerializeToString()\
                               .decode('UTF-8', 'ignore').strip() if x.isprintable())
            if not "BiasAddAdd" in attr_fused_ops:
                continue

            new_node = node_def_pb2.NodeDef()

            new_node.op = quantized_node_op
            new_node.name = requantize_node_name
            for _, value in enumerate(quantized_node.input):
                new_node.input.append(value)

            new_node.input.append(requested_output_min_name)
            new_node.input.append(requested_output_max_name)
            if 'T1' in quantized_node.attr:
                new_node.attr["T1"].CopyFrom(quantized_node.attr['T1'])
            if 'T2' in quantized_node.attr:
                new_node.attr["T2"].CopyFrom(quantized_node.attr['T2'])
            if 'U' in quantized_node.attr:
                new_node.attr["U"].CopyFrom(quantized_node.attr["U"])
            if 'transpose_b' in quantized_node.attr:
                new_node.attr["transpose_b"].CopyFrom(quantized_node.attr['transpose_b'])
            if 'transpose_a' in quantized_node.attr:
                new_node.attr["transpose_a"].CopyFrom(quantized_node.attr['transpose_a'])
            if 'input_quant_mode' in quantized_node.attr:
                new_node.attr["input_quant_mode"].CopyFrom(quantized_node.attr["input_quant_mode"])
            if 'output_quant_mode' in quantized_node.attr:
                new_node.attr["output_quant_mode"].CopyFrom(quantized_node.attr["output_quant_mode"])

            top_node_name = Helper.node_name_from_input(quantized_node.input[0])
            max_filter_node = None
            min_filter_node = None

            # MatMul + BiasAdd + Add
            # The Min and Max of non-const weight node are from QuantizeV2's output, not valid nodes.
            # Add check here for excluding this case.
            if ":2" not in new_node.input[7]:
                max_filter_node = self.graph_info[new_node.input[7]].node
            if ":1" not in new_node.input[6]:
                min_filter_node = self.graph_info[new_node.input[6]].node
            last_node = self.graph_info[new_node.input[0]].node
            weight_node = self.graph_info[Helper.node_name_from_input(new_node.input[1])].node
            bias_node = self.graph_info[Helper.node_name_from_input(new_node.input[2])].node
            if not last_node.op == 'QuantizedConcatV2':
                max_input_node = self.graph_info[last_node.input[-1]].node
                min_input_node = self.graph_info[last_node.input[-2]].node
            
            type_bias = float32_type
            if not last_node.op == 'QuantizedConcatV2' and max_input_node.op == 'Enter': # pragma: no cover
                min_input_parent_name = Helper.node_name_from_input(min_input_node.input[0])
                max_input_parent_name = Helper.node_name_from_input(max_input_node.input[0])
                min_input_parent_node = self.graph_info[min_input_parent_name].node
                max_input_parent_node = self.graph_info[max_input_parent_name].node
                if min_input_parent_node.op != 'Const' or max_input_parent_node.op != 'Const':
                    continue
                min_input_node = min_input_parent_node
                max_input_node = max_input_parent_node
            if max_filter_node and min_filter_node and max_filter_node.op == 'Enter': # pragma: no cover
                min_filter_parent_name = Helper.node_name_from_input(min_filter_node.input[0])
                max_filter_parent_name = Helper.node_name_from_input(max_filter_node.input[0])
                min_filter_parent_node = self.graph_info[min_filter_parent_name].node
                max_filter_parent_node = self.graph_info[max_filter_parent_name].node
                if min_filter_parent_node.op != 'Const' or max_filter_parent_node.op != 'Const':
                    continue
                min_filter_node = min_filter_parent_node
                max_filter_node = max_filter_parent_node
            if weight_node.op == 'Enter': # pragma: no cover
                weight_parent_name = Helper.node_name_from_input(weight_node.input[0])
                weight_parent_node = self.graph_info[weight_parent_name].node
                if weight_parent_node.op != 'Const':
                    continue
                weight_node = weight_parent_node
            bias_enter_node = None
            if bias_node.op == 'Enter': # pragma: no cover
                bias_enter_node = bias_node
                bias_parent_name = Helper.node_name_from_input(bias_node.input[0])
                bias_parent_node = self.graph_info[bias_parent_name].node
                if bias_parent_node.op != 'Const':
                    continue
                bias_node = bias_parent_node

            if max_filter_node and min_filter_node and max_filter_node.op == 'Const' \
              and weight_node.op == 'Const' and not last_node.op == 'QuantizedConcatV2':
                min_input_value = (min_input_node.attr['value'].tensor.float_val)[0]
                max_input_value = (max_input_node.attr['value'].tensor.float_val)[0]
                max_filter_value = (max_filter_node.attr['value'].tensor.float_val)[0]
                min_filter_value = (min_filter_node.attr['value'].tensor.float_val)[0]

                weights_tensor = tensor_util.MakeNdarray(weight_node.attr['value'].tensor)
                bias_tensor = tensor_util.MakeNdarray(bias_node.attr['value'].tensor)
                is_min_first = bool(quantized_node.attr['input_quant_mode'].s == b'MIN_FIRST')
                input_range = max_input_value - min_input_value if is_min_first else max(
                        abs(max_input_value), abs(min_input_value))

                if -self.eps <= input_range <= self.eps:
                    input_range += self.eps

                if -self.eps <= max_input_value - min_input_value <= self.eps:
                    max_input_value += self.eps
                int32_bias = Helper.generate_int32_bias_for_matmul(bias_tensor, weights_tensor,
                                                        input_range, max_input_value,
                                                        min_input_value,
                                                        max_filter_value, min_filter_value)

                bias_node.attr['dtype'].CopyFrom(
                    attr_value_pb2.AttrValue(
                        type=float32_type if self.device == 'gpu' else qint32_type))
                bias_node.attr['value'].CopyFrom(
                    attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
                        bias_tensor if self.device == 'gpu' else int32_bias, dtypes.
                        float32 if self.device == 'gpu' else dtypes.int32, bias_tensor.shape)))

                bias_node.attr['value'].tensor.dtype = float32_type \
                                        if self.device == 'gpu' else qint32_type
                type_bias = float32_type if self.device == 'gpu' else qint32_type
                if bias_enter_node:
                    bias_enter_node.attr["T"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type \
                                                if self.device == 'gpu' else qint32_type))
            else:
                type_bias = float32_type

            new_node.attr["Tbias"].CopyFrom(attr_value_pb2.AttrValue(type=type_bias))

            Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Add', b'Dequantize'])
            Helper.set_attr_type_list(new_node, 'Thost_inputs', [
                uint8_type, int8_type, type_bias, float32_type, float32_type, float32_type, 
                float32_type, float32_type, float32_type, float32_type
                ])

            Helper.set_attr_type_list(new_node, 'Thost_outputs', [float32_type])
            new_node.attr["Tout"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))

            self.graph_analyzer.remove_node(requantize_node_name)

            if self.graph_info[deq_node_name].outputs:
                self.graph_analyzer.replace_single_node(
                    new_node, [top_node_name], quantized_node_name,
                    self.graph_info[deq_node_name].outputs, deq_node_name)
                self.graph_analyzer.remove_node(deq_node_name)
            else:
                self.graph_analyzer.remove_node(deq_node_name)

                new_node.name = deq_node_name
                self.graph_analyzer.replace_single_node(
                    new_node, [top_node_name], quantized_node_name, [], deq_node_name)

            self.graph_analyzer.remove_node(quantized_node_name)

        return self.graph_analyzer.dump_graph()

class FuseMatMulRequantizeNewAPITransformer(GraphRewriterBase):
    """Fuse newAPI Quantized MatMul Op with the successor Requantize Op.
    """
    def __init__(self, model, device='cpu'):
        super().__init__(model)
        self.device = device
        self.graph_analyzer = GraphAnalyzer()
        self.graph_analyzer.graph = self.model
        self.eps = 1e-05
        self.graph_info = self.graph_analyzer.parse_graph()

    def do_transformation(self):
        """Fuse the quantized op with the following requantize op.
        Returns:
            [graphdef]: the optimized graphdef object
        """
        uint8_type = dtypes.quint8.as_datatype_enum
        int8_type = dtypes.qint8.as_datatype_enum
        float32_type = dtypes.float32.as_datatype_enum
        qint32_type = dtypes.qint32.as_datatype_enum

        target_nodes = self.graph_analyzer.query_fusion_pattern_nodes(
            [["_QuantizedMatMul"], ['Requantize']])
        for i in target_nodes:
            quantized_node_name = i[0]
            quantized_node = self.graph_info[quantized_node_name].node
            requantize_node_name = i[1]
            requantize_node = self.graph_info[requantize_node_name].node
            requested_output_min_name = requantize_node.input[3]
            requested_output_max_name = requantize_node.input[4]

            quantized_node_op = i[-1][0]
            attr_fused_ops = ''.join(x for x in quantized_node.attr["fused_ops"].SerializeToString()\
                               .decode('UTF-8', 'ignore').strip() if x.isprintable())
            # "Requantize"
            # "BiasAdd", "Requantize"
            # "BiasAdd", "Activation", "Requantize"
            if "BiasAddAdd" in attr_fused_ops:
                continue

            new_node = node_def_pb2.NodeDef()

            new_node.op = quantized_node_op
            new_node.name = requantize_node_name
            for _, value in enumerate(quantized_node.input):
                new_node.input.append(value)
            new_node.input.append(requested_output_min_name)
            new_node.input.append(requested_output_max_name)

            if 'transpose_b' in quantized_node.attr:
                new_node.attr["transpose_b"].CopyFrom(quantized_node.attr['transpose_b'])
            if 'transpose_a' in quantized_node.attr:
                new_node.attr["transpose_a"].CopyFrom(quantized_node.attr['transpose_a'])
            if 'T1' in quantized_node.attr:
                new_node.attr["T1"].CopyFrom(quantized_node.attr['T1'])
            if 'T2' in quantized_node.attr:
                new_node.attr["T2"].CopyFrom(quantized_node.attr['T2'])
            if 'U' in quantized_node.attr:
                new_node.attr["U"].CopyFrom(quantized_node.attr["U"])
            if 'input_quant_mode' in quantized_node.attr:
                new_node.attr["input_quant_mode"].CopyFrom(quantized_node.attr["input_quant_mode"])
            if 'output_quant_mode' in quantized_node.attr:
                new_node.attr["output_quant_mode"].CopyFrom(quantized_node.attr["output_quant_mode"])

            parent_node_name = Helper.node_name_from_input(quantized_node.input[0])
            max_filter_node = None
            min_filter_node = None
            # The Min and Max of non-const weight node are from QuantizeV2's output, not valid nodes.
            # Add check here for excluding this case.
            if ":2" not in new_node.input[6]:
                max_filter_node = self.graph_info[new_node.input[6]].node
            if ":1" not in new_node.input[5]:
                min_filter_node = self.graph_info[new_node.input[5]].node
            last_node = self.graph_info[new_node.input[0]].node
            is_min_first = bool(quantized_node.attr['input_quant_mode'].s == b'MIN_FIRST')
            weight_node = self.graph_info[new_node.input[1]].node
            bias_node = None
            if "BiasAdd" in attr_fused_ops:
                bias_node = self.graph_info[new_node.input[2]].node
            if not last_node.op == 'QuantizedConcatV2':
                max_input_node = self.graph_info[last_node.input[-1]].node
                min_input_node = self.graph_info[last_node.input[-2]].node

            if not last_node.op == 'QuantizedConcatV2' and max_input_node.op == 'Enter': # pragma: no cover
                min_input_parent_name = Helper.node_name_from_input(min_input_node.input[0])
                max_input_parent_name = Helper.node_name_from_input(max_input_node.input[0])
                min_input_parent_node = self.graph_info[min_input_parent_name].node
                max_input_parent_node = self.graph_info[max_input_parent_name].node
                if min_input_parent_node.op != 'Const' or max_input_parent_node.op != 'Const':
                    continue
                min_input_node = min_input_parent_node
                max_input_node = max_input_parent_node
            if max_filter_node and min_filter_node and max_filter_node.op == 'Enter': # pragma: no cover
                min_filter_parent_name = Helper.node_name_from_input(min_filter_node.input[0])
                max_filter_parent_name = Helper.node_name_from_input(max_filter_node.input[0])
                min_filter_parent_node = self.graph_info[min_filter_parent_name].node
                max_filter_parent_node = self.graph_info[max_filter_parent_name].node
                if min_filter_parent_node.op != 'Const' or max_filter_parent_node.op != 'Const':
                    continue
                min_filter_node = min_filter_parent_node
                max_filter_node = max_filter_parent_node
            if weight_node.op == 'Enter': # pragma: no cover
                weight_parent_name = Helper.node_name_from_input(weight_node.input[0])
                weight_parent_node = self.graph_info[weight_parent_name].node
                if weight_parent_node.op != 'Const':
                    continue
                weight_node = weight_parent_node
            bias_enter_node = None
            if bias_node and bias_node.op == 'Enter': # pragma: no cover
                bias_enter_node = bias_node
                bias_parent_name = Helper.node_name_from_input(bias_node.input[0])
                bias_parent_node = self.graph_info[bias_parent_name].node
                if bias_parent_node.op != 'Const':
                    continue
                bias_node = bias_parent_node

            if bias_node and (last_node.op.find('_QuantizedMatMul') != -1 or last_node.op.find('QuantizeV2') != -1 \
               and max_filter_node and min_filter_node):
                min_input_value = (min_input_node.attr['value'].tensor.float_val)[0]
                max_input_value = (max_input_node.attr['value'].tensor.float_val)[0]
                max_filter_value = (max_filter_node.attr['value'].tensor.float_val)[0]
                min_filter_value = (min_filter_node.attr['value'].tensor.float_val)[0]

                weights_tensor = tensor_util.MakeNdarray(weight_node.attr['value'].tensor)
                bias_tensor = tensor_util.MakeNdarray(bias_node.attr['value'].tensor)

                input_range = max_input_value - min_input_value if is_min_first else max(
                        abs(max_input_value), abs(min_input_value))
                if -self.eps <= input_range <= self.eps:
                    input_range += self.eps

                if -self.eps <= max_input_value - min_input_value <= self.eps:
                    max_input_value += self.eps
                int32_bias = Helper.generate_int32_bias_for_matmul(bias_tensor, weights_tensor,
                                                                   input_range,
                                                                   max_input_value,
                                                                   min_input_value,
                                                                   max_filter_value,
                                                                   min_filter_value)
                bias_node.attr['dtype'].CopyFrom(
                    attr_value_pb2.AttrValue(
                        type=float32_type if self.device == 'gpu' else qint32_type))
                bias_node.attr['value'].CopyFrom(
                    attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
                        bias_tensor if self.device == 'gpu' else int32_bias, dtypes.
                        float32 if self.device == 'gpu' else dtypes.int32, bias_tensor.shape)))

                bias_node.attr['value'].tensor.dtype = float32_type \
                                        if self.device == 'gpu' else qint32_type
                if bias_enter_node:
                    bias_enter_node.attr["T"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type \
                                                if self.device == 'gpu' else qint32_type))
                new_node.attr["Tbias"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type \
                                                if self.device == 'gpu' else qint32_type))

                deq_node_name = self.graph_info[requantize_node_name].outputs[0]
                deq_node = self.graph_info[deq_node_name].node
                deq_node.attr['T'].CopyFrom(attr_value_pb2.AttrValue(type=uint8_type))

                Helper.set_attr_type_list(new_node, "Thost_inputs", [
                                        uint8_type, int8_type,
                                        float32_type if self.device == 'gpu' else qint32_type,
                                        float32_type, float32_type, float32_type,
                                        float32_type, float32_type, float32_type])
            else:
                new_node.attr["Tbias"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))
                deq_node_name = self.graph_info[requantize_node_name].outputs[0]
                deq_node = self.graph_info[deq_node_name].node
                deq_node.attr['T'].CopyFrom(attr_value_pb2.AttrValue(type=uint8_type))
                if bias_node:
                    Helper.set_attr_type_list(new_node, "Thost_inputs", [
                                            uint8_type, int8_type, float32_type,
                                            float32_type, float32_type, float32_type,
                                            float32_type, float32_type, float32_type])
                else:
                    Helper.set_attr_type_list(new_node, "Thost_inputs", [
                                            uint8_type, int8_type,
                                            float32_type, float32_type, float32_type,
                                            float32_type, float32_type, float32_type])

            Helper.set_attr_type_list(new_node, 'Thost_outputs', [
                                      uint8_type, float32_type, float32_type])

            if "GeluApproximate" in  attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', \
                            [b'BiasAdd', b'GeluApproximate', b'Requantize'])
            elif "GeluExact" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'GeluExact', b'Requantize'])
            elif "Elu" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Elu', b'Requantize'])
            elif "LeakyRelu" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'LeakyRelu', b'Requantize'])
            elif "Relu6" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Relu6', b'Requantize'])
            elif "Tanh" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Tanh', b'Requantize'])
            elif "Sigmoid" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Sigmoid', b'Requantize'])
            elif "Relu" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Relu', b'Requantize'])
            elif "BiasAdd" in attr_fused_ops:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Requantize'])
            else:
                Helper.set_attr_string_list(new_node, 'fused_ops', [b'Requantize'])
            new_node.attr["Tout"].CopyFrom(attr_value_pb2.AttrValue(type=uint8_type))

            self.graph_analyzer.replace_single_node(
                new_node, [parent_node_name], quantized_node_name,
                [self.graph_info[requantize_node_name].outputs[0]], requantize_node_name)
            self.graph_analyzer.remove_node(quantized_node_name)

        return self.graph_analyzer.dump_graph()
