| # Copyright (c) Meta Platforms, Inc. and affiliates. |
| # All rights reserved. |
| # |
| # This source code is licensed under the BSD-style license found in the |
| # LICENSE file in the root directory of this source tree. |
| |
| from __future__ import annotations |
| |
| from typing import cast, Optional, Union |
| |
| import torch |
| from executorch.backends.xnnpack._passes.tag_implicit_q_dq_pass import ( |
| TagImplicitQDqPass, |
| ) |
| from executorch.backends.xnnpack.utils.quant_utils import ( |
| extract_qdq_affine_op_args_for_decomposed_ops, |
| is_affine_qdq, |
| is_dequant, |
| is_dynamic_qdq, |
| is_per_channel, |
| is_per_channel_group, |
| is_quant, |
| ) |
| from executorch.backends.xnnpack.utils.utils import ( |
| check_or_raise, |
| get_param_tensor, |
| is_param_node, |
| ) |
| from executorch.exir.dialects._ops import ops as exir_ops |
| from torch.export import ExportedProgram |
| |
| |
| class QuantParams: |
| """ |
| QuantParams class, to represent the paramaters and meta data needed |
| to quantize a tensor. The metadata can technically all be encapsulated |
| within the quant torch.fx.Node, however, there are some cases in which |
| nodes which are meant to be quantized for XNNPACK are not quantized |
| in PyTorch IR, specifically bias nodes. In this case, we can still build |
| quantizer class to serialize the quantized attributes needed for XNNPACK. |
| |
| Attributes: |
| per_channel: Whether this quantization is per channel or per tensor |
| q_input: node that is the input to this quantization |
| scale: tensor or float that is used as the quantization scale |
| zp: tensor or float that is used as the quantization zero point |
| axis: used for per_channel quantizaiton, representing the axis |
| dtype: dtype of the type being quantized to |
| qmin: quantization minimum |
| qmax: quantization maximum |
| is_output: whether this is an output node or not |
| is_input: whether this is an input node or not |
| """ |
| |
| def __init__( |
| self, |
| per_channel: bool, |
| q_input: torch.fx.Node, |
| scale: Union[torch.Tensor, float], |
| zp: Union[torch.Tensor, float], |
| axis: int, |
| dtype: torch.dtype, |
| qmax: int, |
| qmin: int, |
| is_output: bool, |
| is_input: bool, |
| is_dynamic: bool = False, |
| num_nonbatch_dims: int = 1, |
| group_size: int = 0, |
| ) -> None: |
| self.per_channel = per_channel |
| self.q_input = q_input |
| self.scale = scale |
| self.zp = zp |
| self.axis = axis |
| self.dtype = dtype |
| self.qmax = qmax |
| self.qmin = qmin |
| self.is_output = is_output |
| self.is_input = is_input |
| self.is_dynamic = is_dynamic |
| self.num_nonbatch_dims = num_nonbatch_dims |
| self.is_qc4w = ( |
| self.per_channel |
| and not self.is_dynamic |
| and self.qmin == -8 |
| and self.qmax == 7 |
| and self.dtype == torch.int8 |
| ) |
| |
| # Groupwise quantization for weight |
| self.per_channel_group = False |
| self.group_size = group_size |
| if self.group_size > 0: |
| assert ( |
| self.per_channel is True |
| ), "Only per channel quantization supports groupwise quantization" |
| assert ( |
| cast(torch.Tensor, scale).ndim == 2 |
| ), "Scale must be 2D for per channel groupwise quant" |
| self.per_channel_group = True |
| assert group_size > 0, "Group size must be greater than 0" |
| self.is_per_channel_group = self.per_channel and self.group_size > 0 |
| |
| def quantize_tensor(self, tensor: torch.Tensor) -> torch.Tensor: |
| # Do nothing if already quantized by the Quantizer |
| if tensor.dtype == self.dtype: |
| return tensor |
| |
| if self.per_channel: |
| assert ( |
| self.per_channel_group is False |
| ), f"Not expecting per channel group quantization, got q dtype: {self.dtype}, tensor.dtype {tensor.dtype}" |
| assert ( |
| tensor.shape[self.axis] == cast(torch.Tensor, self.scale).shape[0] |
| ), f"Invalid size of per channel quantization scales, axis: {self.axis}, scale size: {self.scale.shape}, tensor shape: {tensor.shape}" |
| |
| assert ( |
| tensor.shape[self.axis] == cast(torch.Tensor, self.zp).shape[0] |
| ), f"Invalid size of per channel quantization zero-points, axis: {self.axis}, zp size: {self.zp.shape}, tensor shape: {tensor.shape}" |
| |
| # Assuming folded quant weights |
| # TODO Add support for unfolded weights |
| assert not self.is_qc4w, "Not expecting QC4W per channel tensor" |
| |
| return exir_ops.edge.quantized_decomposed.quantize_per_channel.default( |
| tensor, self.scale, self.zp, self.axis, self.qmin, self.qmax, self.dtype |
| ) |
| else: |
| return exir_ops.edge.quantized_decomposed.quantize_per_tensor.default( |
| tensor, self.scale, self.zp, self.qmin, self.qmax, self.dtype |
| ) |
| |
| @classmethod |
| def _from_dynamic_input_node(cls, quant_node: torch.fx.Node) -> QuantParams: |
| q_input = quant_node.args[0] # fp32 input |
| assert isinstance(q_input, torch.fx.Node) |
| # TODO - materialize this from the quant_node scale count and val shape |
| num_nonbatch_dims = 1 |
| |
| return cls( |
| per_channel=False, # True is not valid |
| q_input=q_input, |
| scale=0.0, # no need |
| zp=0.0, # no need |
| axis=0, # no need |
| dtype=torch.float32, # will be quantized at runtime |
| qmax=0, # no need |
| qmin=0, # no need |
| is_output=False, |
| is_input=q_input.op == "placeholder", |
| is_dynamic=True, |
| num_nonbatch_dims=num_nonbatch_dims, |
| ) |
| |
| @classmethod |
| def from_q_dq_node( |
| cls, quant_node: torch.fx.Node, ep: Optional[ExportedProgram] = None |
| ) -> QuantParams: |
| check_or_raise( |
| is_quant(quant_node) or is_dequant(quant_node), |
| f"building quantizer from q/dq node but was given node:{quant_node}", |
| ) |
| q_input = quant_node.all_input_nodes[0] |
| |
| # TODO: Use presence of choose_qparam node to determine if this is a dynamic quantization |
| if is_dynamic_qdq(quant_node): |
| return cls._from_dynamic_input_node(quant_node) |
| |
| per_channel = is_per_channel(quant_node) |
| |
| _groupwise = is_per_channel_group(quant_node) |
| quant_node_args = quant_node.args |
| if _groupwise and is_affine_qdq(quant_node): |
| quant_node_args = extract_qdq_affine_op_args_for_decomposed_ops(quant_node) |
| |
| scale = quant_node_args[1] |
| zp = quant_node_args[2] |
| axis = 0 |
| if per_channel: |
| assert isinstance(scale, torch.fx.Node) and isinstance(scale.target, str) |
| assert isinstance(zp, torch.fx.Node) and isinstance(zp.target, str) |
| assert ( |
| ep is not None |
| ), "ExportedProgram must be provided to extract per channel params" |
| |
| def _get_tensor(node): |
| param = get_param_tensor(ep, node) |
| assert param is not None, f"Expected to find param tensor for {node}" |
| return cast(torch.Tensor, param) |
| |
| scale = _get_tensor(scale) |
| zp = _get_tensor(zp) |
| axis = cast(int, quant_node_args[3]) |
| |
| if _groupwise: |
| scale_tensor = cast(torch.Tensor, scale) |
| if scale_tensor.ndim == 1: |
| scale_tensor = scale_tensor.reshape(-1, 1) |
| zp = zp.reshape(-1, 1) |
| scale = scale_tensor |
| |
| assert ( |
| scale_tensor.ndim == 2 |
| ), "Weight scale must be 2D for per_channel_group [de]quant node, got {scale.ndim}D" |
| axis = 0 # axis is ignored for groupwise quantization |
| |
| check_or_raise( |
| bool( |
| quant_node_args[-1] != torch.uint8 |
| or quant_node_args[-1] != torch.quint8 |
| ), |
| "XNNPACK does not support unsigned quantization", |
| ) |
| |
| if _groupwise: |
| _ = quant_node_args[-1] # output dtype - not used |
| group_size = cast(int, quant_node_args[-2]) |
| dtype = cast(torch.dtype, quant_node_args[-3]) |
| qmax = cast(int, quant_node_args[-4]) |
| qmin = cast(int, quant_node_args[-5]) |
| else: |
| group_size = 0 |
| dtype = cast(torch.dtype, quant_node_args[-1]) |
| qmax = cast(int, quant_node_args[-2]) |
| qmin = cast(int, quant_node_args[-3]) |
| |
| is_output = any( |
| user_node.op == "output" for user_node in quant_node.users.keys() |
| ) |
| is_input = q_input.op == "placeholder" |
| return cls( |
| per_channel, |
| q_input, |
| scale, |
| zp, |
| axis, |
| dtype, |
| qmax, |
| qmin, |
| is_output, |
| is_input, |
| group_size=group_size, |
| ) |
| |
| @classmethod |
| def from_weights( |
| cls, tensor_node: torch.fx.Node, ep: Optional[ExportedProgram] = None |
| ) -> Optional[QuantParams]: |
| if not is_dequant(tensor_node): |
| return None |
| |
| # source node for quant params |
| src = tensor_node |
| |
| # is input of dq is q? |
| dq_input = src.all_input_nodes[0] |
| if is_quant(dq_input): |
| src = dq_input |
| |
| # replace this with pointing to the actual weight value. |
| # if no one else uses this weight value then take it out of the toplevel module |
| check_or_raise( |
| src.all_input_nodes[0].op in ["get_attr", "placeholder"], |
| f"q->dq->permute_copy not derived from static weight, input to the q or dq (for folded quant) node: {src.all_input_nodes[0]}", |
| ) |
| |
| return cls.from_q_dq_node(src, ep) |
| |
| @classmethod |
| def from_inputs( |
| cls, tensor_node: torch.fx.Node, ep: ExportedProgram |
| ) -> Optional[QuantParams]: |
| # tensor_node is quantized if it is produced by a dequant node |
| if is_dequant(tensor_node) and TagImplicitQDqPass.is_tagged_as_implicit_q_dq( |
| tensor_node |
| ): |
| dq_input = cast(torch.fx.Node, tensor_node.args[0]) |
| if is_quant(dq_input): |
| q_input = cast(torch.fx.Node, dq_input.args[0]) |
| if is_param_node(ep, q_input): |
| return cls.from_q_dq_node(dq_input) |
| return cls.from_q_dq_node(tensor_node) |
| |
| return None |
| |
| @classmethod |
| def from_outputs(cls, tensor_node: torch.fx.Node) -> Optional[QuantParams]: |
| # tensor_node can also be quantized if it is used as in q -> dq |
| if len(tensor_node.users) == 1: |
| q = list(tensor_node.users.keys())[0] |
| # Check if user is a q node |
| if is_quant(q) and TagImplicitQDqPass.is_tagged_as_implicit_q_dq(q): |
| return cls.from_q_dq_node(q) |
| |
| return None |
| |
| @classmethod |
| def from_bias( |
| cls, |
| bias: torch.fx.Node, |
| weight_quantizer: Optional[QuantParams], |
| input_quantizer: Optional[QuantParams], |
| ) -> Optional[QuantParams]: |
| if weight_quantizer is None or input_quantizer is None: |
| check_or_raise( |
| weight_quantizer is None and input_quantizer is None, |
| "Weight and Input should both be quantized", |
| ) |
| return None |
| |
| if input_quantizer.is_dynamic: |
| # No need to quantize bias for dyanamic quantization |
| return None |
| |
| check_or_raise( |
| not input_quantizer.per_channel, |
| "Input can not be quantized per channel", |
| ) |
| |
| # Only per_tensor quantization is supported for input here |
| check_or_raise( |
| isinstance(input_quantizer.scale, float), |
| f"q_input scale should be float, but got {input_quantizer.scale}", |
| ) |
| return cls( |
| per_channel=weight_quantizer.per_channel, |
| q_input=bias, |
| scale=weight_quantizer.scale * cast(float, input_quantizer.scale), |
| zp=weight_quantizer.zp * 0, |
| axis=0, # not using weight_quantizer.axis because bias is always of shape [out_channels] i.e. 1D |
| dtype=torch.int32, |
| qmin=-(2**31), |
| qmax=(2**31) - 1, |
| is_output=False, |
| is_input=False, |
| ) |
