|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +from dataclasses import replace |
| 8 | +from unittest.mock import Mock |
| 9 | + |
| 10 | +import pytest |
| 11 | +import torch |
| 12 | +import tosa_serializer as ts |
| 13 | +from executorch.backends.arm._passes.quant_args import QuantArgs |
| 14 | +from executorch.backends.arm.operators.ops_quant_utils import add_input_weight_zp_consts |
| 15 | +from executorch.backends.arm.quantizer.arm_quantizer import ( |
| 16 | + get_symmetric_quantization_config, |
| 17 | +) |
| 18 | +from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig |
| 19 | +from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineINT |
| 20 | +from executorch.backends.test.harness.stages import StageType |
| 21 | +from torchao.quantization.pt2e.quantizer import ( |
| 22 | + DerivedQuantizationSpec, |
| 23 | + QuantizationSpec, |
| 24 | +) |
| 25 | + |
| 26 | + |
| 27 | +def _affine_config(per_channel: bool) -> QuantizationConfig: |
| 28 | + config = get_symmetric_quantization_config(is_per_channel=per_channel) |
| 29 | + weight = config.weight |
| 30 | + assert isinstance(weight, QuantizationSpec) |
| 31 | + return replace( |
| 32 | + config, |
| 33 | + weight=replace( |
| 34 | + weight, |
| 35 | + qscheme=( |
| 36 | + torch.per_channel_affine if per_channel else torch.per_tensor_affine |
| 37 | + ), |
| 38 | + quant_min=-128, |
| 39 | + ), |
| 40 | + ) |
| 41 | + |
| 42 | + |
| 43 | +@pytest.mark.parametrize("per_channel", [False, True]) |
| 44 | +def test_affine_weight_specs_and_derived_channel_bias(per_channel: bool) -> None: |
| 45 | + config = _affine_config(per_channel) |
| 46 | + assert config.get_weight_qspec() is config.weight |
| 47 | + graph = torch.fx.Graph() |
| 48 | + node = graph.call_function( |
| 49 | + torch.ops.aten.linear.default, |
| 50 | + ( |
| 51 | + graph.placeholder("input"), |
| 52 | + graph.placeholder("weight"), |
| 53 | + graph.placeholder("bias"), |
| 54 | + ), |
| 55 | + ) |
| 56 | + bias = config.get_bias_qspec(node) |
| 57 | + assert isinstance(bias, DerivedQuantizationSpec) |
| 58 | + assert bias.ch_axis == (0 if per_channel else None) |
| 59 | + |
| 60 | + |
| 61 | +def test_equal_channel_weight_zero_points_serialize_as_one_scalar() -> None: |
| 62 | + graph = torch.fx.Graph() |
| 63 | + node = graph.placeholder("conv") |
| 64 | + node.meta["input_qparams"] = { |
| 65 | + 0: QuantArgs(0.1, 0, -127, 127, torch.int8), |
| 66 | + 1: QuantArgs([0.01, 0.02], [-28, -28], -128, 127, torch.int8, per_channel=True), |
| 67 | + } |
| 68 | + inputs = [Mock(dtype=ts.DType.INT8), Mock(dtype=ts.DType.INT8)] |
| 69 | + serializer = Mock() |
| 70 | + input_const, weight_const = Mock(), Mock() |
| 71 | + input_const.name, weight_const.name = "input_zp", "weight_zp" |
| 72 | + serializer.addConst.side_effect = [input_const, weight_const] |
| 73 | + add_input_weight_zp_consts(serializer, node, inputs, "conv") |
| 74 | + assert serializer.addConst.call_args_list[1].args == ([1], ts.DType.INT8, [-28]) |
| 75 | + node.meta["input_qparams"][1] = QuantArgs( |
| 76 | + [0.01, 0.02], [-28, -27], -128, 127, torch.int8, per_channel=True |
| 77 | + ) |
| 78 | + with pytest.raises(ValueError, match="common weight zero point"): |
| 79 | + add_input_weight_zp_consts(serializer, node, inputs, "conv") |
| 80 | + |
| 81 | + |
| 82 | +def _affine_model( |
| 83 | + operator: str, conflicting_zero_points: bool = False |
| 84 | +) -> tuple[torch.nn.Module, tuple[torch.Tensor]]: |
| 85 | + model: torch.nn.Linear | torch.nn.Conv2d | torch.nn.ConvTranspose2d |
| 86 | + if operator == "linear": |
| 87 | + model = torch.nn.Linear(8, 3) |
| 88 | + inputs = (torch.linspace(-1, 1, 8).reshape(1, 8),) |
| 89 | + else: |
| 90 | + model = ( |
| 91 | + torch.nn.ConvTranspose2d(2, 3, 2) |
| 92 | + if operator == "conv_transpose2d" |
| 93 | + else torch.nn.Conv2d(2, 3, 2) |
| 94 | + ) |
| 95 | + inputs = (torch.linspace(-1, 1, 72).reshape(1, 2, 6, 6),) |
| 96 | + axis = 1 if operator == "conv_transpose2d" else 0 |
| 97 | + channel_shape = model.weight.movedim(axis, 0).shape |
| 98 | + base = torch.linspace(-0.25, 1, model.weight.numel() // 3).reshape( |
| 99 | + channel_shape[1:] |
| 100 | + ) |
| 101 | + weights = torch.stack([base * scale for scale in (1, 2, 3)]) |
| 102 | + if conflicting_zero_points: |
| 103 | + weights[0] += 0.25 |
| 104 | + with torch.no_grad(): |
| 105 | + model.weight.copy_(weights.movedim(0, axis)) |
| 106 | + assert model.bias is not None |
| 107 | + model.bias.copy_(torch.tensor([0.1, 0.2, 0.3])) |
| 108 | + return model.eval(), inputs |
| 109 | + |
| 110 | + |
| 111 | +def _check_affine_qparams( |
| 112 | + pipeline: TosaPipelineINT[tuple[torch.Tensor]], per_channel: bool |
| 113 | +) -> None: |
| 114 | + graph_module = pipeline.tester.get_artifact(StageType.QUANTIZE) |
| 115 | + op = ( |
| 116 | + torch.ops.quantized_decomposed.dequantize_per_channel.default |
| 117 | + if per_channel |
| 118 | + else torch.ops.quantized_decomposed.dequantize_per_tensor.default |
| 119 | + ) |
| 120 | + weights = [ |
| 121 | + node |
| 122 | + for node in graph_module.graph.nodes |
| 123 | + if node.target == op and node.args[-1] == torch.int8 |
| 124 | + ] |
| 125 | + if per_channel: |
| 126 | + assert len(weights) == 1 |
| 127 | + scales = graph_module.get_buffer(str(weights[0].args[1].target)) |
| 128 | + zero_points = graph_module.get_buffer(str(weights[0].args[2].target)) |
| 129 | + assert scales.numel() == 3 |
| 130 | + assert torch.unique(scales).numel() == 3 |
| 131 | + assert torch.unique(zero_points).numel() == 1 |
| 132 | + assert zero_points[0] != 0 |
| 133 | + biases = [ |
| 134 | + node |
| 135 | + for node in graph_module.graph.nodes |
| 136 | + if node.target == op and node.args[-1] == torch.int32 |
| 137 | + ] |
| 138 | + assert len(biases) == 1 |
| 139 | + assert biases[0].args[3] == 0 |
| 140 | + assert graph_module.get_buffer(str(biases[0].args[1].target)).numel() == 3 |
| 141 | + else: |
| 142 | + weight = next(node for node in weights if node.args[0].op == "get_attr") |
| 143 | + assert weight.args[2] != 0 |
| 144 | + |
| 145 | + |
| 146 | +@pytest.mark.parametrize("operator", ["linear", "conv2d", "conv_transpose2d"]) |
| 147 | +@pytest.mark.parametrize("per_channel", [False, True]) |
| 148 | +def test_affine_weights_tosa_pipeline(operator: str, per_channel: bool) -> None: |
| 149 | + model, inputs = _affine_model(operator) |
| 150 | + pipeline = TosaPipelineINT[tuple[torch.Tensor]]( |
| 151 | + model, |
| 152 | + inputs, |
| 153 | + aten_op=f"torch.ops.aten.{operator}.{'input' if operator == 'conv_transpose2d' else 'default'}", |
| 154 | + per_channel_quantization=per_channel, |
| 155 | + ) |
| 156 | + pipeline.quantizer.set_global(_affine_config(per_channel)) |
| 157 | + pipeline.add_stage_after("quantize", _check_affine_qparams, pipeline, per_channel) |
| 158 | + pipeline.run() |
| 159 | + assert pipeline.has_stage("run_method_and_compare_outputs") |
| 160 | + |
| 161 | + |
| 162 | +def test_conflicting_affine_weight_zero_points_rejected_by_pipeline() -> None: |
| 163 | + model, inputs = _affine_model("conv2d", conflicting_zero_points=True) |
| 164 | + pipeline = TosaPipelineINT[tuple[torch.Tensor]]( |
| 165 | + model, inputs, aten_op="torch.ops.aten.conv2d.default" |
| 166 | + ) |
| 167 | + pipeline.quantizer.set_global(_affine_config(True)) |
| 168 | + with pytest.raises(ValueError, match="common weight zero point"): |
| 169 | + pipeline.run() |
0 commit comments