diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 2bf0892a7e9..e5103635f89 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -31,6 +31,7 @@ fbcode_target(_kind = runtime.python_library, "activation_fusion_pass.py", "aten_to_cortex_m_pass.py", "clamp_hardswish_pass.py", + "compute_quantized_constant_ops_pass.py", "cortex_m_pass.py", "cortex_m_pass_manager.py", "decompose_hardswish_pass.py", diff --git a/backends/cortex_m/passes/compute_quantized_constant_ops_pass.py b/backends/cortex_m/passes/compute_quantized_constant_ops_pass.py new file mode 100644 index 00000000000..b7d19812a5e --- /dev/null +++ b/backends/cortex_m/passes/compute_quantized_constant_ops_pass.py @@ -0,0 +1,34 @@ +# 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. + +import torch +from executorch.backends.arm._passes import ComputeConstantOpsAOTPass +from executorch.exir.dialects._ops import ops as exir_ops + + +class ComputeQuantizedConstantOpsPass(ComputeConstantOpsAOTPass): + """Give the quantized constant ops their int8 values. + + Folding removes the Q node after these ops but leaves them producing float, + so the int8 kernels they feed would read float data. A ``full`` stays a + runtime fill with its fill value quantized, so it adds no constant data; the + other ops are computed in float and stored as quantized constants. Runs + after ``ConvertFullLikeToFullPass``, so a ``full_like`` arrives as ``full``. + """ + + def compute_node_aot(self, node: torch.fx.Node) -> bool: + output_qparams = node.meta.get("output_qparams") + if not output_qparams: + return False + if node.target == exir_ops.edge.aten.full.default: + qparams = output_qparams[0] + fill_value = node.args[1] + assert isinstance(fill_value, (int, float)) + fill = qparams.quantize_value(torch.tensor(float(fill_value))) + node.args = (node.args[0], int(fill.item())) + node.kwargs = {**node.kwargs, "dtype": qparams.dtype} + return True + return super().compute_node_aot(node) diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index c5e475c58a7..5127d030fc3 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -10,6 +10,7 @@ import torch from executorch.backends.arm._passes import ( + ConvertFullLikeToFullPass, DeduplicateGetAttrPass, FoldAndAnnotateQParamsPass, ScalarsToAttributePass, @@ -48,6 +49,7 @@ from .activation_fusion_pass import ActivationFusionPass from .aten_to_cortex_m_pass import AtenToCortexMPass from .clamp_hardswish_pass import ClampHardswishPass +from .compute_quantized_constant_ops_pass import ComputeQuantizedConstantOpsPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass from .decompose_sdpa_pass import DecomposeSDPAPass @@ -149,6 +151,8 @@ class CortexMPassManager(ExportedProgramPassManager): FoldAndAnnotateQParamsPass, ActivationFusionPass, QuantizedClampActivationPass, + ConvertFullLikeToFullPass, + ComputeQuantizedConstantOpsPass, DecomposeHardswishPass, AtenToCortexMPass, # Retrace only after quantized ATen ops have become Cortex-M kernels. @@ -165,6 +169,8 @@ class CortexMPassManager(ExportedProgramPassManager): ReplaceScalarWithTensorArgPass, ActivationFusionPass, QuantizedClampActivationPass, + ConvertFullLikeToFullPass, + ComputeQuantizedConstantOpsPass, DecomposeHardswishPass, ConvertConv1dToConv2dPass, CortexMReplaceOpsWithChannelsLastVariants, diff --git a/backends/cortex_m/test/misc/test_quantization.py b/backends/cortex_m/test/misc/test_quantization.py index 6f27d1a6c28..4b3c8e971ae 100644 --- a/backends/cortex_m/test/misc/test_quantization.py +++ b/backends/cortex_m/test/misc/test_quantization.py @@ -6,7 +6,7 @@ import torch from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor -from executorch.backends.arm.test.common import parametrize, xfail_type +from executorch.backends.arm.test.common import parametrize from executorch.backends.cortex_m.test.tester import ( CortexMTester, McuTestCase, @@ -358,12 +358,8 @@ def forward(self, x): ), } -xfails: dict[str, xfail_type] = { - "surrounded_quantized_op_constant": "Numerical error since the add is forced to have non-correct qparams.", -} - -@parametrize("test_case", test_cases, xfails=xfails) +@parametrize("test_case", test_cases) def test_shared_qspec_quantizer(test_case): """ Test that ops which does not change dynamic range are able to use int8 portable kernels. diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index 689d117e7b0..d29248e5647 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -546,14 +546,126 @@ def test_convs_around_a_reshape_behind_lut_activation_lower_to_cmsis_nn(): tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) -def test_explicit_layout_rejects_float_input_to_int8_conv(): +def test_explicit_layout_lowers_conv_on_a_quantized_constant(): + inputs = (torch.randn(1, 8, 16),) model = ThenConv1d(lambda x: torch.full_like(x, 0.5, dtype=torch.float32)) - tester = CortexMTester(model.eval(), (torch.randn(1, 8, 16),)) + tester = _run_explicit_layout_passes(CortexMTester(model.eval(), inputs)) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() - with pytest.raises(Exception) as caught: - _run_explicit_layout_passes(tester) + assert _count(program, exir_ops.edge.aten.full_like.default) == 0 + assert _count(program, exir_ops.edge.aten.full.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) - error = caught.value - while error.__cause__ is not None: - error = error.__cause__ - assert "must move an int8 tensor" in str(error) + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +@pytest.mark.parametrize( + "constant,constant_op,shape", + [ + ( + lambda: torch.arange(-1.0, 1.0, 0.125), + exir_ops.edge.aten.arange.start_step, + (16,), + ), + ( + lambda: torch.scalar_tensor(0.5), + torch.ops.aten.scalar_tensor.default, + (), + ), + ], + ids=["arange", "scalar_tensor"], +) +def test_quantized_constant_is_computed_before_quantizing( + use_explicit_layout, constant, constant_op, shape +): + torch.manual_seed(0) + inputs = (torch.randn(shape),) + tester = CortexMTester(Apply(lambda x: x + constant()).eval(), inputs) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + assert _count(program, constant_op) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +@pytest.mark.parametrize( + "fill", + [lambda x: torch.full((1, 8, 16), 0.5), lambda x: torch.full_like(x, 0.5)], + ids=["full", "full_like"], +) +def test_quantized_fill_stays_a_runtime_int8_full(use_explicit_layout, fill): + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = CortexMTester(Apply(lambda x: x + fill(x)).eval(), inputs) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + (full,) = program.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.full.default + ) + assert full.meta["val"].dtype == torch.int8 + assert _count(program, exir_ops.edge.aten.full_like.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +def test_clamp_with_constant_tensor_bounds_quantizes_them_once(use_explicit_layout): + torch.manual_seed(0) + inputs = (torch.randn(2, 3, 4, 5),) + tester = CortexMTester( + Apply( + lambda x: torch.clamp(x, torch.full_like(x, -1.0), torch.full_like(x, 1.0)) + ).eval(), + inputs, + ) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + (clamp,) = program.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.clamp.default + ) + qparams = clamp.meta["output_qparams"][0] + expected = tuple( + int(qparams.quantize_value(torch.tensor(bound)).item()) for bound in (-1.0, 1.0) + ) + assert tuple(clamp.args[1:]) == expected + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +def test_unquantized_full_like_keeps_int8_after_float_promoting_op(use_explicit_layout): + inputs = (torch.randn(1, 8, 16),) + tester = CortexMTester( + Apply( + lambda x: ( + torch.sigmoid(x).permute(0, 2, 1), + torch.full_like(x, 2, dtype=torch.int32), + ) + ).eval(), + inputs, + ) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout))) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 1 + assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0 + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +def test_unquantized_constant_is_not_precomputed(use_explicit_layout): + inputs = (torch.randn(1, 8, 16),) + tester = CortexMTester( + Apply(lambda x: x + torch.full((1, 8, 16), 0.5)).eval(), inputs + ) + tester.export().to_edge() + program = _run_passes(tester, use_explicit_layout) + + assert _count(program, exir_ops.edge.aten.full.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs)