TFNewAPI support Quantized Matmul, BatchMatmul (#1116)

lvliang-intel · web-flow · commit 70a85d4afb5f · 2022-08-04T10:47:57.000+08:00
diff --git a/neural_compressor/adaptor/inteltensorflow.yaml b/neural_compressor/adaptor/inteltensorflow.yaml
@@ -279,8 +279,11 @@
         'Dequantize + MatMul + BiasAdd + Relu + QuantizeV2',
         'Dequantize + MatMul + BiasAdd + QuantizeV2',
         'Dequantize + MatMul + Relu + QuantizeV2',
-        'Dequantize + BatchMatMulV2',
-        'Dequantize + BatchMatMulV2 + Mul + Add',
+        'Dequantize + BatchMatMulV2 + Mul + QuantizeV2',
+        'Dequantize + BatchMatMulV2 + Add + QuantizeV2',
+        'Dequantize + BatchMatMulV2 + AddV2 + QuantizeV2',
+        'Dequantize + BatchMatMulV2 + Mul + Add + QuantizeV2',
+        'Dequantize + BatchMatMulV2 + Mul + AddV2 + QuantizeV2',
         'Dequantize + Conv3D + AddV2 + AddV2 + Relu + QuantizeV2',
         'Dequantize + Conv3D + Add + Relu + QuantizeV2',
         'Dequantize + Conv3D + AddV2 + Relu + QuantizeV2',
diff --git a/neural_compressor/adaptor/tensorflow.py b/neural_compressor/adaptor/tensorflow.py
@@ -599,7 +599,8 @@ def _dump_model_op_stats(self, model_graphdef):
         int8_op_prefix_list = ['QuantizedConv2D', '_QuantizedConv3D', 'QuantizedDepthwise',
                                'QuantizedMaxPool', 'QuantizedAvgPool',
                                'QuantizedConcatV2', 'QuantizedMatMul',
-                               '_QuantizedFusedBatchNorm']
+                               '_QuantizedFusedBatchNorm', '_QuantizedMatMul',
+                               '_QuantizedBatchMatMul']
         from tensorflow.python.framework import dtypes
 
         res = {}
@@ -620,6 +621,8 @@ def _dump_model_op_stats(self, model_graphdef):
                     origin_op_type = 'FusedBatchNormV3'
                 if origin_op_type == 'Depthwise':
                     origin_op_type = 'DepthwiseConv2dNative'
+                if origin_op_type == 'BatchMatMul':
+                    origin_op_type = 'BatchMatMulV2'
                 res[origin_op_type]['INT8'] += 1
 
             if i.op in fp32_op_list:
diff --git a/neural_compressor/adaptor/tf_utils/graph_converter.py b/neural_compressor/adaptor/tf_utils/graph_converter.py
@@ -56,6 +56,8 @@
 from .graph_rewriter.int8.fuse_conv_requantize import FuseConvRequantizeTransformer
 from .graph_rewriter.int8.fuse_matmul_requantize import FuseMatMulRequantizeTransformer
 from .graph_rewriter.int8.fuse_matmul_requantize import FuseMatMulRequantizeDequantizeTransformer
+from .graph_rewriter.int8.fuse_matmul_requantize import FuseMatMulRequantizeNewAPITransformer
+from .graph_rewriter.int8.fuse_matmul_requantize import FuseMatMulRequantizeDequantizeNewAPITransformer
 from .graph_rewriter.int8.scale_propagation import ScaleProPagationTransformer
 from .graph_rewriter.bf16.bf16_convert import BF16Convert
 from .graph_rewriter.int8.post_quantized_op_cse import PostCseOptimizer
@@ -547,7 +549,7 @@ def _freeze_requantization_ranges(self, additional_data=None):
         self.scale_info.update(requant_min_max)
 
         self._tmp_graph_def = QuantizedRNNConverter(
-            self._tmp_graph_def, self._calibration_data, self._rnn_details).do_transformation()
+            self._tmp_graph_def, self._calibration_data, self._rnn_details, self.new_api).do_transformation()
 
         if 'scale_propagation_max_pooling' in self.recipes and \
                 self.recipes['scale_propagation_max_pooling']:
@@ -570,18 +572,18 @@ def _fuse_requantize_with_fused_quantized_node(self):
 
         if not self.fake_quant:
             # TODO Use MatMul and BatchMatMul new API
-            #if self.qdq_enabled:
-            #    self._tmp_graph_def = FuseMatMulRequantizeNewAPITransformer(
-            #        self._tmp_graph_def).do_transformation()
-            #
-            #    self._tmp_graph_def = FuseMatMulRequantizeDequantizeNewAPITransformer(
-            #        self._tmp_graph_def).do_transformation()
-            #else:
-            self._tmp_graph_def = FuseMatMulRequantizeTransformer(
-                self._tmp_graph_def).do_transformation()
+            if self.qdq_enabled:
+                self._tmp_graph_def = FuseMatMulRequantizeNewAPITransformer(
+                    self._tmp_graph_def).do_transformation()
+            
+                self._tmp_graph_def = FuseMatMulRequantizeDequantizeNewAPITransformer(
+                    self._tmp_graph_def).do_transformation()
+            else:
+                self._tmp_graph_def = FuseMatMulRequantizeTransformer(
+                            self._tmp_graph_def).do_transformation()
 
-            self._tmp_graph_def = FuseMatMulRequantizeDequantizeTransformer(
-                self._tmp_graph_def).do_transformation()
+                self._tmp_graph_def = FuseMatMulRequantizeDequantizeTransformer(
+                            self._tmp_graph_def).do_transformation()
 
         self._tmp_graph_def = StripUnusedNodesOptimizer(
             self._tmp_graph_def,
diff --git a/neural_compressor/adaptor/tf_utils/graph_rewriter/bf16/bf16_convert.py b/neural_compressor/adaptor/tf_utils/graph_rewriter/bf16/bf16_convert.py
@@ -129,11 +129,10 @@ def _bf16_convert(self, bf16_node_name):
             return
         else:
             self.converted_ops.append(bf16_node.name)
-
+        
         inputs_dt, outputs_dt = self._dtype(bf16_node)
         inputs_dt_val, outputs_dt_val = self._dtype_val(bf16_node)
         allowed_dt_val = self._allowed_dtype_val(bf16_node)
-
         for index, input_name in enumerate(bf16_node.input):
             if input_name.startswith('^'):
                 continue
@@ -142,7 +141,6 @@ def _bf16_convert(self, bf16_node_name):
                 input_name)]
             input_node = input_detail.node
             input_node_outputs = input_detail.outputs
-
             if inputs_dt[index] in allowed_dt_val and \
                                         dtypes.bfloat16.as_datatype_enum not in allowed_dt_val[inputs_dt[index]]:
                 continue
@@ -239,6 +237,10 @@ def _model_bf16_convert(self):
             if bf16_node_name not in self.cur_graph.node_name_details:
                 self.bf16_ops.remove(bf16_node_name)
                 continue
+            else:
+                if "fused_ops" in self.cur_graph.node_name_details[bf16_node_name].node.attr:
+                    self.bf16_ops.remove(bf16_node_name)
+                    continue
         for bf16_node_name in set(self.bf16_ops):
             self._bf16_convert(bf16_node_name)
         return self.cur_graph.dump_graph()
diff --git a/neural_compressor/adaptor/tf_utils/graph_rewriter/int8/fuse_matmul_requantize.py b/neural_compressor/adaptor/tf_utils/graph_rewriter/int8/fuse_matmul_requantize.py
@@ -261,8 +261,8 @@ def do_transformation(self):
 
         return self.graph_analyzer.dump_graph()
 
-class FuseMatMulRequantizeDequantizeNewAPITransformer(GraphRewriterBase): # pragma: no cover
-    """Fuse _QuantizedFusedMatMul + Requantize + Dequantize into _QuantizedFusedMatMulAndDequantize.
+class FuseMatMulRequantizeDequantizeNewAPITransformer(GraphRewriterBase):
+    """Fuse _QuantizedMatMul + Requantize + Dequantize into _QuantizedMatMul.
     """
     def __init__(self, model, device='cpu'):
         super().__init__(model)
@@ -275,20 +275,13 @@ def __init__(self, model, device='cpu'):
         self.eps = 1e-5
 
     def do_transformation(self):
-        fuse_pattern = [["_QuantizedFusedMatMul"], ['Requantize'], ['Dequantize'], ('Softmax',)]
+        fuse_pattern = [["_QuantizedMatMul"], ['Requantize'], ['Dequantize'], ('Softmax',)]
 
         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:
-                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]
@@ -299,26 +292,30 @@ def do_transformation(self):
 
             new_node = node_def_pb2.NodeDef()
 
-            new_node.op = quantized_node_op + "AndDequantize"
+            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 'num_args' in quantized_node.attr:
-                new_node.attr["num_args"].CopyFrom(quantized_node.attr['num_args'])
+            if 'Tbias' in quantized_node.attr:
+                new_node.attr["Tbias"].CopyFrom(quantized_node.attr['Tbias'])
             if 'fused_ops' in quantized_node.attr:
                 new_node.attr["fused_ops"].CopyFrom(quantized_node.attr["fused_ops"])
-
+            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"])
+            if 'Thost_inputs' in quantized_node.attr:
+                new_node.attr["Thost_inputs"].CopyFrom(quantized_node.attr["Thost_inputs"])
+            Helper.set_attr_type_list(new_node, 'Thost_outputs', [dtypes.float32.as_datatype_enum])
+            Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Dequantize'])
             top_node_name = Helper.node_name_from_input(quantized_node.input[0])
             float32_type = dtypes.float32.as_datatype_enum
-            new_node.attr["Targs"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))
-            new_node.attr["Toutput"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))
+            new_node.attr["Tout"].CopyFrom(attr_value_pb2.AttrValue(type=float32_type))
 
             self.graph_analyzer.remove_node(requantize_node_name)
 
@@ -338,7 +335,7 @@ def do_transformation(self):
 
         return self.graph_analyzer.dump_graph()
 
-class FuseMatMulRequantizeNewAPITransformer(GraphRewriterBase): # pragma: no cover
+class FuseMatMulRequantizeNewAPITransformer(GraphRewriterBase):
     """Fuse newAPI Quantized MatMul Op with the successor Requantize Op.
     """
     def __init__(self, model, device='cpu'):
@@ -358,7 +355,7 @@ def do_transformation(self):
 
         while True:
             target_nodes = self.graph_analyzer.query_fusion_pattern_nodes(
-                [["_QuantizedFusedMatMul"], ['Requantize']])
+                [["_QuantizedMatMul"], ['Requantize']])
             if len(target_nodes) == 0:
                 break
 
@@ -377,23 +374,41 @@ def do_transformation(self):
 
             new_node = node_def_pb2.NodeDef()
 
-            new_node.op = quantized_node_op + "AndRequantize"
+            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 'num_args' in quantized_node.attr:
-                new_node.attr["num_args"].CopyFrom(quantized_node.attr["num_args"])
-            if 'Targs' in quantized_node.attr:
-                new_node.attr["Targs"].CopyFrom(quantized_node.attr["Targs"])
-            if 'fused_ops' in quantized_node.attr:
-                new_node.attr["fused_ops"].CopyFrom(quantized_node.attr["fused_ops"])
-            new_node.attr["Toutput"].CopyFrom(attr_value_pb2.AttrValue(type=uint8_type))
+            if 'Tbias' in quantized_node.attr:
+                new_node.attr["Tbias"].CopyFrom(quantized_node.attr["Targs"])
+            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"])
+            Helper.set_attr_type_list(new_node, "Thost_inputs", [
+                                      dtypes.quint8.as_datatype_enum,
+                                      dtypes.qint8.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum])
+            Helper.set_attr_type_list(new_node, 'Thost_outputs', [
+                                      dtypes.quint8.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum,
+                                      dtypes.float32.as_datatype_enum])
+            Helper.set_attr_string_list(new_node, 'fused_ops', [b'BiasAdd', b'Relu', b'Requantize'])
+            new_node.attr["Tout"].CopyFrom(attr_value_pb2.AttrValue(type=uint8_type))
 
             parent_node_name = Helper.node_name_from_input(quantized_node.input[0])
             self.graph_analyzer.replace_single_node(
diff --git a/neural_compressor/adaptor/tf_utils/graph_rewriter/int8/rnn_convert.py b/neural_compressor/adaptor/tf_utils/graph_rewriter/int8/rnn_convert.py
@@ -30,10 +30,11 @@
 
 
 class QuantizedRNNConverter(GraphRewriterBase):
-    def __init__(self, model, calibration_data, rnn_details):
+    def __init__(self, model, calibration_data, rnn_details, new_api=False):
         super().__init__(model)
         self.calibration_data = calibration_data
         self.rnn_details = rnn_details
+        self.new_api = new_api
 
     @dump_elapsed_time("Pass QuantizedRNNConverter")
     def do_transformation(self):
@@ -207,22 +208,51 @@ def do_transformation(self):
 
             quantized_matmul_input.append(enter_min_node.name)
             quantized_matmul_input.append(enter_max_node.name)
-            quantized_matmul_with_bias_node = Helper.create_node(
-                'QuantizedMatMulWithBias', i[0] + '_quantized_mat_mul', quantized_matmul_input)
+            if self.new_api:
+                quantized_matmul_with_bias_node = Helper.create_node(
+                    '_QuantizedMatMul', i[0] + '_quantized_mat_mul', quantized_matmul_input)
+            else:
+                quantized_matmul_with_bias_node = Helper.create_node(
+                    'QuantizedMatMulWithBias', i[0] + '_quantized_mat_mul', quantized_matmul_input)
             Helper.set_attr_dtype(
                 quantized_matmul_with_bias_node, 'T1', dtypes.quint8)
             Helper.set_attr_dtype(
                 quantized_matmul_with_bias_node, 'T2', dtypes.qint8)
             Helper.set_attr_dtype(
                 quantized_matmul_with_bias_node, 'Tbias', dtypes.float32)
-            Helper.set_attr_dtype(
-                quantized_matmul_with_bias_node, 'Toutput', dtypes.qint32)
+            if self.new_api: 
+                Helper.set_attr_dtype(
+                    quantized_matmul_with_bias_node, 'Tout', dtypes.qint32)
+            else:
+                Helper.set_attr_dtype(
+                    quantized_matmul_with_bias_node, 'Toutput', dtypes.qint32)
             Helper.set_attr_bool(
                 quantized_matmul_with_bias_node, 'transpose_a', False)
             Helper.set_attr_bool(
                 quantized_matmul_with_bias_node, 'transpose_b', False)
-            Helper.set_attr_string(
-                quantized_matmul_with_bias_node, 'input_quant_mode', b"MIN_FIRST")
+            if self.new_api:
+                Helper.set_attr_string(
+                    quantized_matmul_with_bias_node, 'input_quant_mode', b"SCALED")
+                Helper.set_attr_string(
+                    quantized_matmul_with_bias_node, 'output_quant_mode', b"SCALED")
+                Helper.set_attr_string_list(quantized_matmul_with_bias_node, 'fused_ops', [b'BiasAdd'])
+                Helper.set_attr_type_list(quantized_matmul_with_bias_node, 'Thost_inputs', [
+                        dtypes.quint8.as_datatype_enum,
+                        dtypes.qint8.as_datatype_enum,
+                        dtypes.float32.as_datatype_enum,
+                        dtypes.float32.as_datatype_enum,
+                        dtypes.float32.as_datatype_enum,
+                        dtypes.float32.as_datatype_enum,
+                        dtypes.float32.as_datatype_enum
+                     ])
+                Helper.set_attr_type_list(quantized_matmul_with_bias_node, 'Thost_outputs', [
+                                          dtypes.qint32.as_datatype_enum,
+                                          dtypes.float32.as_datatype_enum,
+                                          dtypes.float32.as_datatype_enum])
+            else:
+                Helper.set_attr_string(
+                    quantized_matmul_with_bias_node, 'input_quant_mode', b"MIN_FIRST")
+
             g.add_node(quantized_matmul_with_bias_node,
                        quantize_node.name, [bias_node.name])
 
diff --git a/neural_compressor/adaptor/tf_utils/graph_rewriter/qdq/insert_qdq_pattern.py b/neural_compressor/adaptor/tf_utils/graph_rewriter/qdq/insert_qdq_pattern.py
@@ -157,9 +157,10 @@ def _insert_qdq_pattern_for_common_ops(self, original_node, is_asymmetric):
             if each_input_name[0] == '^':
                 continue
 
-            if self.node_name_mapping[original_node.name].op == "MatMul" or \
-               self.node_name_mapping[original_node.name].op == "BatchMatMulV2":
+            if self.node_name_mapping[original_node.name].op == "MatMul":
                 dtype = dtypes.quint8
+            elif self.node_name_mapping[original_node.name].op == "BatchMatMulV2":
+                dtype = dtypes.qint8
             else:
                 input_node_name = Helper.node_name_from_input(each_input_name)
                 if input_node_name in self.graph_info:
diff --git a/neural_compressor/adaptor/tf_utils/quantize_graph/qdq/fuse_qdq_matmul.py b/neural_compressor/adaptor/tf_utils/quantize_graph/qdq/fuse_qdq_matmul.py
diff --git a/test/tfnewapi/test_tensorflow_graph_qdq_matmul_fusion.py b/test/tfnewapi/test_tensorflow_graph_qdq_matmul_fusion.py
