Upload 379 files

610d9c2 verified about 3 hours ago

8 kB

import itertools
import logging

import onnx

from ..quant_utils import (
    FLOAT8_TYPES,
    TENSOR_NAME_QUANT_SUFFIX,
    QuantizedValue,
    QuantizedValueType,
    find_by_name,
    get_mul_node,
)
from .base_operator import QuantOperatorBase
from .qdq_base_operator import QDQOperatorBase


class QOpMatMul(QuantOperatorBase):
    def __init__(self, onnx_quantizer, onnx_node):
        super().__init__(onnx_quantizer, onnx_node)

    def should_quantize(self):
        if not self.quantizer.should_quantize_node(self.node):
            logging.debug(f"Ignore MatMul {self.node.name}]")
            return False

        if (not self.quantizer.is_float_tensor(self.node.input[1])) and (
            not self.quantizer.is_float_tensor(self.node.input[0])
        ):
            logging.info(f"Ignore MatMul due to non float inputs {self.node.name}]")
            return False

        # do not quantize non-constant B matrices for matmul
        if self.quantizer.q_matmul_const_b_only:
            if not self.quantizer.find_initializer_in_path(self.node.input[1]):
                logging.info(f"Ignore MatMul due to non constant B: {self.quantizer.graph_scope}[{self.node.name}]")
                return False
        return True


"""
    Used when quantize mode is QuantizationMode.IntegerOps.
"""


class MatMulInteger(QOpMatMul):
    def __init__(self, onnx_quantizer, onnx_node):
        super().__init__(onnx_quantizer, onnx_node)

    def quantize(self):
        node = self.node
        assert node.op_type == "MatMul"
        #  Get Quantized from both activation(input[0]) and weight(input[1])
        (
            quantized_input_names,
            zero_point_names,
            scale_names,
            nodes,
        ) = self.quantizer.quantize_activation(node, [0])

        (
            quantized_input_names_weight,
            zero_point_names_weight,
            scale_names_weight,
            nodes_weight,
        ) = self.quantizer.quantize_weight(node, [1], reduce_range=True, op_level_per_channel=True)
        quantized_input_names.extend(quantized_input_names_weight)
        zero_point_names.extend(zero_point_names_weight)
        scale_names.extend(scale_names_weight)
        nodes.extend(nodes_weight)

        matmul_integer_output = node.output[0] + "_output_quantized"
        matmul_integer_name = node.name + "_quant" if node.name else ""
        matmul_integer_node = onnx.helper.make_node(
            "MatMulInteger",
            quantized_input_names + zero_point_names,
            [matmul_integer_output],
            matmul_integer_name,
        )
        nodes.append(matmul_integer_node)

        # Add cast operation to cast matmulInteger output to float.
        cast_op_output = matmul_integer_output + "_cast_output"
        otype = self.quantizer.get_tensor_type(node.output[0], mandatory=True)
        cast_node = onnx.helper.make_node(
            "Cast",
            [matmul_integer_output],
            [cast_op_output],
            matmul_integer_output + "_cast",
            to=otype,
        )
        nodes.append(cast_node)

        # Add mul operation to multiply scales of two inputs.
        assert len(scale_names) == 2
        scales_mul_op = (
            matmul_integer_name + "_scales_mul"
            if matmul_integer_name
            else scale_names[0] + "_" + scale_names[1] + "_mul"
        )

        scales_mul_node = find_by_name(scales_mul_op, self.quantizer.new_nodes)
        if scales_mul_node is None:
            scales_mul_node = get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op)
            nodes.append(scales_mul_node)

        scales_mul_op_output = scales_mul_node.output[0]

        # Add mul operation to multiply mul_scales_op result with output of MatMulInteger
        # and make the output of this node the same as output of original matmul node.
        output_scale_mul_op = ""
        if matmul_integer_name:
            output_scale_mul_op = matmul_integer_name + "_output_scale_mul"
        nodes.append(
            get_mul_node(
                [cast_op_output, scales_mul_op_output],
                node.output[0],
                output_scale_mul_op,
            )
        )
        self.quantizer.new_nodes += nodes


"""
    Used when quantize mode is QuantizationMode.QLinearOps
"""


class QLinearMatMul(QOpMatMul):
    def __init__(self, onnx_quantizer, onnx_node):
        super().__init__(onnx_quantizer, onnx_node)

    def quantize(self):
        node = self.node
        assert node.op_type == "MatMul"
        #  Get Quantized from both activation(input[0]) and weight(input[1])
        (
            quantized_input_names,
            zero_point_names,
            scale_names,
            nodes,
        ) = self.quantizer.quantize_activation(node, [0])

        (
            quantized_input_names_weight,
            zero_point_names_weight,
            scale_names_weight,
            nodes_weight,
        ) = self.quantizer.quantize_weight(node, [1], reduce_range=True, op_level_per_channel=True)
        quantized_input_names.extend(quantized_input_names_weight)
        zero_point_names.extend(zero_point_names_weight)
        scale_names.extend(scale_names_weight)

        nodes.extend(nodes_weight)
        (
            data_found,
            output_scale_name,
            output_zp_name,
            _,
            _,
        ) = self.quantizer._get_quantization_params(node.output[0])
        if not data_found or quantized_input_names is None:
            return super().quantize()

        qlinear_matmul_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
        qlinear_matmul_name = node.name + "_quant" if node.name else ""

        qlinear_matmul_inputs = []
        # Input 0
        qlinear_matmul_inputs.append(quantized_input_names[0])
        qlinear_matmul_inputs.append(scale_names[0])
        qlinear_matmul_inputs.append(zero_point_names[0])
        # Input 1
        qlinear_matmul_inputs.append(quantized_input_names[1])
        qlinear_matmul_inputs.append(scale_names[1])
        qlinear_matmul_inputs.append(zero_point_names[1])
        # Output quantization parameter
        qlinear_matmul_inputs.append(output_scale_name)
        qlinear_matmul_inputs.append(output_zp_name)

        domain = "com.microsoft" if self.quantizer.weight_qType in FLOAT8_TYPES else ""
        qlinear_matmul_node = onnx.helper.make_node(
            "QLinearMatMul",
            qlinear_matmul_inputs,
            [qlinear_matmul_output],
            qlinear_matmul_name,
            domain=domain,
        )
        nodes.append(qlinear_matmul_node)

        # Create an entry for this quantized value
        q_output = QuantizedValue(
            node.output[0],
            qlinear_matmul_output,
            output_scale_name,
            output_zp_name,
            QuantizedValueType.Input,
        )
        self.quantizer.quantized_value_map[node.output[0]] = q_output

        self.quantizer.new_nodes += nodes


class QDQMatMul(QDQOperatorBase):
    def __init__(self, onnx_quantizer, onnx_node):
        super().__init__(onnx_quantizer, onnx_node)

    def quantize(self):
        node = self.node
        assert node.op_type == "MatMul"

        if self.disable_qdq_for_node_output:
            nodes_to_iterate = node.input
        else:
            nodes_to_iterate = itertools.chain(node.input, node.output)

        for tensor_name in nodes_to_iterate:
            if find_by_name(tensor_name, self.quantizer.model.initializer()):
                is_per_channel, channel_axis = self.quantizer.is_tensor_per_channel(
                    tensor_name, default_axis=1, op_type=node.op_type
                )
                if is_per_channel:
                    self.quantizer.quantize_weight_tensor_per_channel(tensor_name, channel_axis)
                else:
                    self.quantizer.quantize_weight_tensor(tensor_name)
            else:
                self.quantizer.quantize_activation_tensor(tensor_name)

	import itertools
	import logging

	import onnx

	from ..quant_utils import (
	FLOAT8_TYPES,
	TENSOR_NAME_QUANT_SUFFIX,
	QuantizedValue,
	QuantizedValueType,
	find_by_name,
	get_mul_node,
	)
	from .base_operator import QuantOperatorBase
	from .qdq_base_operator import QDQOperatorBase


	class QOpMatMul(QuantOperatorBase):
	def __init__(self, onnx_quantizer, onnx_node):
	super().__init__(onnx_quantizer, onnx_node)

	def should_quantize(self):
	if not self.quantizer.should_quantize_node(self.node):
	logging.debug(f"Ignore MatMul {self.node.name}]")
	return False

	if (not self.quantizer.is_float_tensor(self.node.input[1])) and (
	not self.quantizer.is_float_tensor(self.node.input[0])
	):
	logging.info(f"Ignore MatMul due to non float inputs {self.node.name}]")
	return False

	# do not quantize non-constant B matrices for matmul
	if self.quantizer.q_matmul_const_b_only:
	if not self.quantizer.find_initializer_in_path(self.node.input[1]):
	logging.info(f"Ignore MatMul due to non constant B: {self.quantizer.graph_scope}[{self.node.name}]")
	return False
	return True


	"""
	Used when quantize mode is QuantizationMode.IntegerOps.
	"""


	class MatMulInteger(QOpMatMul):
	def __init__(self, onnx_quantizer, onnx_node):
	super().__init__(onnx_quantizer, onnx_node)

	def quantize(self):
	node = self.node
	assert node.op_type == "MatMul"
	# Get Quantized from both activation(input[0]) and weight(input[1])
	(
	quantized_input_names,
	zero_point_names,
	scale_names,
	nodes,
	) = self.quantizer.quantize_activation(node, [0])

	(
	quantized_input_names_weight,
	zero_point_names_weight,
	scale_names_weight,
	nodes_weight,
	) = self.quantizer.quantize_weight(node, [1], reduce_range=True, op_level_per_channel=True)
	quantized_input_names.extend(quantized_input_names_weight)
	zero_point_names.extend(zero_point_names_weight)
	scale_names.extend(scale_names_weight)
	nodes.extend(nodes_weight)

	matmul_integer_output = node.output[0] + "_output_quantized"
	matmul_integer_name = node.name + "_quant" if node.name else ""
	matmul_integer_node = onnx.helper.make_node(
	"MatMulInteger",
	quantized_input_names + zero_point_names,
	[matmul_integer_output],
	matmul_integer_name,
	)
	nodes.append(matmul_integer_node)

	# Add cast operation to cast matmulInteger output to float.
	cast_op_output = matmul_integer_output + "_cast_output"
	otype = self.quantizer.get_tensor_type(node.output[0], mandatory=True)
	cast_node = onnx.helper.make_node(
	"Cast",
	[matmul_integer_output],
	[cast_op_output],
	matmul_integer_output + "_cast",
	to=otype,
	)
	nodes.append(cast_node)

	# Add mul operation to multiply scales of two inputs.
	assert len(scale_names) == 2
	scales_mul_op = (
	matmul_integer_name + "_scales_mul"
	if matmul_integer_name
	else scale_names[0] + "_" + scale_names[1] + "_mul"
	)

	scales_mul_node = find_by_name(scales_mul_op, self.quantizer.new_nodes)
	if scales_mul_node is None:
	scales_mul_node = get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op)
	nodes.append(scales_mul_node)

	scales_mul_op_output = scales_mul_node.output[0]

	# Add mul operation to multiply mul_scales_op result with output of MatMulInteger
	# and make the output of this node the same as output of original matmul node.
	output_scale_mul_op = ""
	if matmul_integer_name:
	output_scale_mul_op = matmul_integer_name + "_output_scale_mul"
	nodes.append(
	get_mul_node(
	[cast_op_output, scales_mul_op_output],
	node.output[0],
	output_scale_mul_op,
	)
	)
	self.quantizer.new_nodes += nodes


	"""
	Used when quantize mode is QuantizationMode.QLinearOps
	"""


	class QLinearMatMul(QOpMatMul):
	def __init__(self, onnx_quantizer, onnx_node):
	super().__init__(onnx_quantizer, onnx_node)

	def quantize(self):
	node = self.node
	assert node.op_type == "MatMul"
	# Get Quantized from both activation(input[0]) and weight(input[1])
	(
	quantized_input_names,
	zero_point_names,
	scale_names,
	nodes,
	) = self.quantizer.quantize_activation(node, [0])

	(
	quantized_input_names_weight,
	zero_point_names_weight,
	scale_names_weight,
	nodes_weight,
	) = self.quantizer.quantize_weight(node, [1], reduce_range=True, op_level_per_channel=True)
	quantized_input_names.extend(quantized_input_names_weight)
	zero_point_names.extend(zero_point_names_weight)
	scale_names.extend(scale_names_weight)

	nodes.extend(nodes_weight)
	(
	data_found,
	output_scale_name,
	output_zp_name,
	_,
	_,
	) = self.quantizer._get_quantization_params(node.output[0])
	if not data_found or quantized_input_names is None:
	return super().quantize()

	qlinear_matmul_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
	qlinear_matmul_name = node.name + "_quant" if node.name else ""

	qlinear_matmul_inputs = []
	# Input 0
	qlinear_matmul_inputs.append(quantized_input_names[0])
	qlinear_matmul_inputs.append(scale_names[0])
	qlinear_matmul_inputs.append(zero_point_names[0])
	# Input 1
	qlinear_matmul_inputs.append(quantized_input_names[1])
	qlinear_matmul_inputs.append(scale_names[1])
	qlinear_matmul_inputs.append(zero_point_names[1])
	# Output quantization parameter
	qlinear_matmul_inputs.append(output_scale_name)
	qlinear_matmul_inputs.append(output_zp_name)

	domain = "com.microsoft" if self.quantizer.weight_qType in FLOAT8_TYPES else ""
	qlinear_matmul_node = onnx.helper.make_node(
	"QLinearMatMul",
	qlinear_matmul_inputs,
	[qlinear_matmul_output],
	qlinear_matmul_name,
	domain=domain,
	)
	nodes.append(qlinear_matmul_node)

	# Create an entry for this quantized value
	q_output = QuantizedValue(
	node.output[0],
	qlinear_matmul_output,
	output_scale_name,
	output_zp_name,
	QuantizedValueType.Input,
	)
	self.quantizer.quantized_value_map[node.output[0]] = q_output

	self.quantizer.new_nodes += nodes


	class QDQMatMul(QDQOperatorBase):
	def __init__(self, onnx_quantizer, onnx_node):
	super().__init__(onnx_quantizer, onnx_node)

	def quantize(self):
	node = self.node
	assert node.op_type == "MatMul"

	if self.disable_qdq_for_node_output:
	nodes_to_iterate = node.input
	else:
	nodes_to_iterate = itertools.chain(node.input, node.output)

	for tensor_name in nodes_to_iterate:
	if find_by_name(tensor_name, self.quantizer.model.initializer()):
	is_per_channel, channel_axis = self.quantizer.is_tensor_per_channel(
	tensor_name, default_axis=1, op_type=node.op_type
	)
	if is_per_channel:
	self.quantizer.quantize_weight_tensor_per_channel(tensor_name, channel_axis)
	else:
	self.quantizer.quantize_weight_tensor(tensor_name)
	else:
	self.quantizer.quantize_activation_tensor(tensor_name)