codekingpro/portable-devtools

1import itertools
2import logging
3
4import onnx
5from onnx import onnx_pb as onnx_proto
6
7from ..quant_utils import TENSOR_NAME_QUANT_SUFFIX, QuantizedValue, QuantizedValueType, find_by_name, get_mul_node
8from .base_operator import QuantOperatorBase
9from .qdq_base_operator import QDQOperatorBase
10
11
12class QOpMatMul(QuantOperatorBase):
13    def __init__(self, onnx_quantizer, onnx_node):
14        super().__init__(onnx_quantizer, onnx_node)
15
16    def should_quantize(self):
17        if not self.quantizer.should_quantize_node(self.node):
18            logging.debug(f"Ignore MatMul {self.node.name}]")
19            return False
20
21        if (not self.quantizer.is_float_tensor(self.node.input[1])) and (
22            not self.quantizer.is_float_tensor(self.node.input[0])
23        ):
24            logging.info(f"Ignore MatMul due to non float inputs {self.node.name}]")
25            return False
26
27        # do not quantize non-constant B matrices for matmul
28        if self.quantizer.q_matmul_const_b_only:
29            if not self.quantizer.find_initializer_in_path(self.node.input[1]):
30                logging.info(f"Ignore MatMul due to non constant B: {self.quantizer.graph_scope}[{self.node.name}]")
31                return False
32        return True
33
34
35"""
36    Used when quantize mode is QuantizationMode.IntegerOps.
37"""
38
39
40class MatMulInteger(QOpMatMul):
41    def __init__(self, onnx_quantizer, onnx_node):
42        super().__init__(onnx_quantizer, onnx_node)
43
44    def quantize(self):
45        node = self.node
46        assert node.op_type == "MatMul"
47        #  Get Quantized from both activation(input[0]) and weight(input[1])
48        (
49            quantized_input_names,
50            zero_point_names,
51            scale_names,
52            nodes,
53        ) = self.quantizer.quantize_activation(node, [0])
54
55        (
56            quantized_input_names_weight,
57            zero_point_names_weight,
58            scale_names_weight,
59            nodes_weight,
60        ) = self.quantizer.quantize_weight(node, [1], reduce_range=True, op_level_per_channel=True)
61        quantized_input_names.extend(quantized_input_names_weight)
62        zero_point_names.extend(zero_point_names_weight)
63        scale_names.extend(scale_names_weight)
64        nodes.extend(nodes_weight)
65
66        matmul_integer_output = node.output[0] + "_output_quantized"
67        matmul_integer_name = node.name + "_quant" if node.name else ""
68        matmul_integer_node = onnx.helper.make_node(
69            "MatMulInteger",
70            quantized_input_names + zero_point_names,
71            [matmul_integer_output],
72            matmul_integer_name,
73        )
74        nodes.append(matmul_integer_node)
75
76        # Add cast operation to cast matmulInteger output to float.
77        cast_op_output = matmul_integer_output + "_cast_output"
78        otype = self.quantizer.get_tensor_type(node.output[0], mandatory=True)
79        cast_node = onnx.helper.make_node(
80            "Cast",
81            [matmul_integer_output],
82            [cast_op_output],
83            matmul_integer_output + "_cast",
84            to=otype,
85        )
86        nodes.append(cast_node)
87
88        # Add mul operation to multiply scales of two inputs.
89        assert len(scale_names) == 2
90        scales_mul_op = (
91            matmul_integer_name + "_scales_mul"
92            if matmul_integer_name
93            else scale_names[0] + "_" + scale_names[1] + "_mul"
94        )
95
96        scales_mul_node = find_by_name(scales_mul_op, self.quantizer.new_nodes)
97        if scales_mul_node is None:
98            scales_mul_node = get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op)
99            nodes.append(scales_mul_node)
100
101        scales_mul_op_output = scales_mul_node.output[0]
102
103        # Add mul operation to multiply mul_scales_op result with output of MatMulInteger
104        # and make the output of this node the same as output of original matmul node.
105        output_scale_mul_op = ""
106        if matmul_integer_name:
107            output_scale_mul_op = matmul_integer_name + "_output_scale_mul"
108        nodes.append(
109            get_mul_node(
110                [cast_op_output, scales_mul_op_output],
111                node.output[0],
112                output_scale_mul_op,
113            )
114        )
115        self.quantizer.new_nodes += nodes
116
117
118"""
119    Used when quantize mode is QuantizationMode.QLinearOps
120"""
121
122
123class QLinearMatMul(QOpMatMul):
124    def __init__(self, onnx_quantizer, onnx_node):
125        super().__init__(onnx_quantizer, onnx_node)
126
127    def quantize(self):
128        node = self.node
129        assert node.op_type == "MatMul"
130        #  Get Quantized from both activation(input[0]) and weight(input[1])
131        (
132            quantized_input_names,
133            zero_point_names,
134            scale_names,
135            nodes,
136        ) = self.quantizer.quantize_activation(node, [0])
137
138        (
139            quantized_input_names_weight,
140            zero_point_names_weight,
141            scale_names_weight,
142            nodes_weight,
143        ) = self.quantizer.quantize_weight(node, [1], reduce_range=True, op_level_per_channel=True)
144        quantized_input_names.extend(quantized_input_names_weight)
145        zero_point_names.extend(zero_point_names_weight)
146        scale_names.extend(scale_names_weight)
147
148        nodes.extend(nodes_weight)
149        (
150            data_found,
151            output_scale_name,
152            output_zp_name,
153            _,
154            _,
155        ) = self.quantizer._get_quantization_params(node.output[0])
156        if not data_found or quantized_input_names is None:
157            return super().quantize()
158
159        qlinear_matmul_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
160        qlinear_matmul_name = node.name + "_quant" if node.name else ""
161
162        qlinear_matmul_inputs = []
163        # Input 0
164        qlinear_matmul_inputs.append(quantized_input_names[0])
165        qlinear_matmul_inputs.append(scale_names[0])
166        qlinear_matmul_inputs.append(zero_point_names[0])
167        # Input 1
168        qlinear_matmul_inputs.append(quantized_input_names[1])
169        qlinear_matmul_inputs.append(scale_names[1])
170        qlinear_matmul_inputs.append(zero_point_names[1])
171        # Output quantization parameter
172        qlinear_matmul_inputs.append(output_scale_name)
173        qlinear_matmul_inputs.append(output_zp_name)
174
175        domain = (
176            "com.microsoft"
177            if self.quantizer.weight_qType
178            in {
179                onnx_proto.TensorProto.FLOAT8E4M3FN,
180                onnx_proto.TensorProto.FLOAT8E4M3FNUZ,
181                onnx_proto.TensorProto.FLOAT8E5M2,
182                onnx_proto.TensorProto.FLOAT8E5M2FNUZ,
183            }
184            else ""
185        )
186        qlinear_matmul_node = onnx.helper.make_node(
187            "QLinearMatMul",
188            qlinear_matmul_inputs,
189            [qlinear_matmul_output],
190            qlinear_matmul_name,
191            domain=domain,
192        )
193        nodes.append(qlinear_matmul_node)
194
195        # Create an entry for this quantized value
196        q_output = QuantizedValue(
197            node.output[0],
198            qlinear_matmul_output,
199            output_scale_name,
200            output_zp_name,
201            QuantizedValueType.Input,
202        )
203        self.quantizer.quantized_value_map[node.output[0]] = q_output
204
205        self.quantizer.new_nodes += nodes
206
207
208class QDQMatMul(QDQOperatorBase):
209    def __init__(self, onnx_quantizer, onnx_node):
210        super().__init__(onnx_quantizer, onnx_node)
211
212    def quantize(self):
213        node = self.node
214        assert node.op_type == "MatMul"
215
216        if self.disable_qdq_for_node_output:
217            nodes_to_iterate = node.input
218        else:
219            nodes_to_iterate = itertools.chain(node.input, node.output)
220
221        for tensor_name in nodes_to_iterate:
222            if find_by_name(tensor_name, self.quantizer.model.initializer()):
223                is_per_channel, channel_axis = self.quantizer.is_tensor_per_channel(
224                    tensor_name, default_axis=1, op_type=node.op_type
225                )
226                if is_per_channel:
227                    self.quantizer.quantize_weight_tensor_per_channel(tensor_name, channel_axis)
228                else:
229                    self.quantizer.quantize_weight_tensor(tensor_name)
230            else:
231                self.quantizer.quantize_activation_tensor(tensor_name)
232