Pattern for quantized linear: dequantize_affine + linear.
| 826 | |
| 827 | @REGISTRY.register_pattern(name="QUANTIZED_LINEAR") |
| 828 | class QuantizedLinearHandler(PatternHandler): |
| 829 | """ |
| 830 | Pattern for quantized linear: dequantize_affine + linear. |
| 831 | """ |
| 832 | |
| 833 | def __init__( |
| 834 | self, |
| 835 | head: Node, |
| 836 | body: List[Node], |
| 837 | qdata: Node, |
| 838 | scale: Node, |
| 839 | zero_point: Node, |
| 840 | group_size: int, |
| 841 | bits: int, |
| 842 | out_dtype: torch.dtype, |
| 843 | ): |
| 844 | super().__init__(head, body) |
| 845 | self.qdata = qdata |
| 846 | self.scale = scale |
| 847 | self.zero_point = zero_point |
| 848 | self.group_size = group_size |
| 849 | self.bits = bits |
| 850 | self.out_dtype = out_dtype |
| 851 | |
| 852 | @classmethod |
| 853 | def maybe_create( |
| 854 | cls, ep: ExportedProgram, head: Node |
| 855 | ) -> Optional["QuantizedLinearHandler"]: |
| 856 | linear_node = head |
| 857 | if not match_target(linear_node, torch.ops.aten.linear.default): |
| 858 | return None |
| 859 | |
| 860 | x, w = linear_node.args[0:2] |
| 861 | dequant_node = w |
| 862 | if not match_target(dequant_node, torch.ops.torchao.dequantize_affine.default): |
| 863 | return None |
| 864 | if not has_single_user(dequant_node): |
| 865 | return None |
| 866 | |
| 867 | parsed = parse_dequant_node(dequant_node) |
| 868 | if parsed is None: |
| 869 | return None |
| 870 | qdata, scale, zero_point, group_size, bits, out_dtype, _quantized_dim = parsed |
| 871 | out_dtype = x.meta["val"].dtype if out_dtype is None else out_dtype |
| 872 | |
| 873 | head = linear_node |
| 874 | body = [dequant_node] |
| 875 | return QuantizedLinearHandler( |
| 876 | head, |
| 877 | body, |
| 878 | qdata=qdata, |
| 879 | scale=scale, |
| 880 | zero_point=zero_point, |
| 881 | group_size=group_size, |
| 882 | bits=bits, |
| 883 | out_dtype=out_dtype, |
| 884 | ) |
| 885 |