diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 2bf0892a7e9..16404979e4d 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -40,6 +40,7 @@ fbcode_target(_kind = runtime.python_library, "fuse_conv_padding_pass.py", "initialize_scratch_buffers_pass.py", "matmul_to_bmm_pass.py", + "propagate_qparams_pass.py", "quantized_clamp_activation_pass.py", "scratch_buffer_sizes.py", ], diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index b9339b838b8..7ae1ecb174a 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -24,6 +24,7 @@ from executorch.backends.cortex_m.utils import ( build_activation_lut, is_foldable_alpha, + lowered_output_dtype, quantize_multiplier_aot, quantize_val, SHIFT_INT8, @@ -1209,8 +1210,7 @@ def _get_minimum_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass - input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype not in (torch.int8, torch.int32): + if lowered_output_dtype(node.args[0]) not in (torch.int8, torch.int32): return None return DialectNodeSpec(exir_ops.edge.cortex_m.minimum.default, node.args) @@ -1220,8 +1220,7 @@ def _get_maximum_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass - input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype != torch.int8: + if lowered_output_dtype(node.args[0]) != torch.int8: return None return DialectNodeSpec(exir_ops.edge.cortex_m.maximum.default, node.args) @@ -1241,8 +1240,12 @@ def _get_permute_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass + if lowered_output_dtype(node.args[0]) != torch.int8: + return None + # Shape, unlike dtype, survives AtenToDialectPass: a substitution keeps the + # replaced node's meta and computes the same function over it. input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype != torch.int8 or not 1 <= input_tensor.dim() <= 4: + if not 1 <= input_tensor.dim() <= 4: return None return _transpose_spec(node, input_tensor) diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 1bd12cbaab0..fd8f32c2a00 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -54,6 +54,7 @@ from .fuse_conv_padding_pass import FuseConvPaddingPass from .initialize_scratch_buffers_pass import InitializeScratchBuffersPass from .matmul_to_bmm_pass import MatmulToBmmPass +from .propagate_qparams_pass import PropagateQParamsPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -139,6 +140,8 @@ class CortexMPassManager(ExportedProgramPassManager): # Move layout copies across pads before singleton permutations become views. RemovePermutesAroundElementwiseOps, CortexMCanonicalizeViewCopyPermutePass, + # After the plumbing has settled, before anything asks what it carries. + PropagateQParamsPass, ValidateCortexMExplicitLayoutPass, ReplaceQuantNodesPass, AtenToCortexMPass, diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py index 659b67d729e..f1063bd8eb6 100644 --- a/backends/cortex_m/passes/explicit_layout_pass.py +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -10,6 +10,7 @@ import torch from executorch.backends.cortex_m.utils import ( coerce_int_pair, + lowered_output_dtype, skips_quantized_max_pool2d, ) from executorch.backends.transforms.canonicalize_view_copy_permute_pass import ( @@ -134,12 +135,15 @@ def call(self, graph_module: GraphModule) -> PassResult: input_value = ( input_node.meta.get("val") if isinstance(input_node, Node) else None ) - if ( - not isinstance(input_value, torch.Tensor) - or input_value.dtype != torch.int8 - ): + input_dtype = ( + lowered_output_dtype(input_node) + if isinstance(input_node, Node) + else None + ) + if not isinstance(input_value, torch.Tensor) or input_dtype != torch.int8: raise RuntimeError( - f"Cortex-M layout copy {node.name} must move an int8 tensor." + f"Cortex-M layout copy {node.name} must move an int8 tensor, " + f"but its input {input_node} is {input_dtype}." ) rank = input_value.dim() if ( diff --git a/backends/cortex_m/passes/propagate_qparams_pass.py b/backends/cortex_m/passes/propagate_qparams_pass.py new file mode 100644 index 00000000000..800b72b4d2a --- /dev/null +++ b/backends/cortex_m/passes/propagate_qparams_pass.py @@ -0,0 +1,82 @@ +# 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 logging + +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass +from torch.fx import GraphModule, Node +from torch.fx.node import Target +from torch.fx.passes.infra.pass_manager import PassResult + +logger: logging.Logger = logging.getLogger(__name__) + +# Ops that rearrange a tensor without changing its element type: a reshape, a +# layout copy, a transpose. Whatever dtype the producer hands them is the dtype +# they hand on. +_TRANSPARENT_TARGETS: frozenset[Target] = frozenset( + ( + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.squeeze_copy.dim, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.transpose_copy.int, + *PERMUTE_COPY_TARGETS, + ) +) + + +class PropagateQParamsPass(ExportPass): + """Give the shape and layout plumbing the qparams of the tensor it moves. + + FoldAndAnnotateQParamsPass annotates the ops that compute. The ops that + only rearrange are left with neither qparams nor a refreshed fake tensor, + so they still carry the fp32 tensor the pre-quantization trace gave them, + and anything asking what a plumbing node produces -- `lowered_output_dtype` + and so every caller gating a substitution on int8, plus + ValidateCortexMExplicitLayoutPass -- is told fp32 about a tensor that is + int8 on device. The symptom is not a wrong answer but a missing one: a + substitution declines, the op stays on a portable kernel, and nothing says + so. + + Rearranging does not change an element type, so a plumbing node's qparams + are its producer's. Annotating them here resolves the question once, in the + graph, rather than at each of the places that ask it. Nodes are visited in + topological order, so a chain of plumbing resolves in a single pass: each + node reads a producer this pass has already annotated. + + Runs before ValidateCortexMExplicitLayoutPass, which is the first consumer, + and late enough that the passes which read and write `output_qparams` + themselves -- ActivationFusionPass, DecomposeHardswishPass -- have already + seen the graph without these annotations. + """ + + def call(self, graph_module: GraphModule) -> PassResult: + modified = False + + for node in graph_module.graph.nodes: + if node.op != "call_function" or node.target not in _TRANSPARENT_TARGETS: + continue + if node.meta.get("output_qparams"): + continue + if not node.args or not isinstance(producer := node.args[0], Node): + continue + qparams = producer.meta.get("output_qparams") + if not qparams: + continue + # The same QuantArgs the producer carries: a rearrangement reads and + # writes one tensor, so there is a single output to describe and it + # is quantized exactly as the input was. + node.meta["output_qparams"] = dict(qparams) + modified = True + logger.debug( + "[PropagateQParams] %s inherits qparams from %s", + node.name, + producer.name, + ) + + return PassResult(graph_module, modified) diff --git a/backends/cortex_m/passes/quantized_clamp_activation_pass.py b/backends/cortex_m/passes/quantized_clamp_activation_pass.py index 429d1ebb445..59f8f1cadb1 100644 --- a/backends/cortex_m/passes/quantized_clamp_activation_pass.py +++ b/backends/cortex_m/passes/quantized_clamp_activation_pass.py @@ -7,11 +7,14 @@ from typing import Any import torch -from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( get_output_qparams, ) -from executorch.backends.cortex_m.utils import get_activation_bounds, quantize_val +from executorch.backends.cortex_m.utils import ( + get_activation_bounds, + lowered_output_dtype, + quantize_val, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass from torch.fx import GraphModule, Node @@ -71,11 +74,7 @@ def _is_quantized_int8_activation(self, node: Node) -> bool: input_node = node.args[0] if len(node.args) > 0 else None if not isinstance(input_node, Node): return False - try: - tensor = get_first_fake_tensor(input_node) - except Exception: - return False - if tensor is None or tensor.dtype != torch.int8: + if lowered_output_dtype(input_node) != torch.int8: return False try: diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index ac48ee03869..c7d60cbdb4e 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -105,6 +105,19 @@ def define_common_targets(is_fbcode = False): ], ) + python_unittest( + name = "test_propagate_qparams", + srcs = ["test_propagate_qparams.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/arm/_passes:passes", + "//executorch/backends/cortex_m/passes:cortex_passes", + "//executorch/backends/cortex_m:utils", + "//executorch/backends/cortex_m/ops:ops", + "//executorch/exir/dialects:lib", + ], + ) + python_pytest( name = "test_attention", srcs = ["ops/test_attention.py"], diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index 30528fff9ed..dbdef44c31c 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -20,6 +20,7 @@ CortexMRunPasses, CortexMTester, ) +from executorch.backends.cortex_m.utils import lowered_output_dtype from executorch.backends.test.harness.stages import Quantize, StageType from executorch.backends.transforms.remove_unused_constants_pass import ( RemoveUnusedConstantsPass, @@ -67,6 +68,87 @@ def forward(self, x): ) +class LutActivationThenConv1d(torch.nn.Module): + """A LUT activation feeding a convolution. + + tanh lowers to ``cortex_m::quantized_activation``, but only in + ``AtenToCortexMPass``, which runs after the layout validator. Until then + it keeps an fp32 fake tensor and passes that dtype on to the convolution, + so validating on ``val`` rejects this graph even though it lowers. + An activation that ``ActivationFusionPass`` folds into the convolution + clamp, such as relu, never exposes the gap. + """ + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + return self.conv(torch.tanh(x)) + + +class LutActivationViewThenConv1d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + return self.conv(torch.tanh(x).view(1, 8, 16)) + + +class LutActivationReluThenConv1d(torch.nn.Module): + """A relu between a LUT activation and a convolution. + + Nothing fuses this relu. ``ActivationFusionPass`` only folds an + activation that follows a convolution, so it reaches + ``QuantizedClampActivationPass``, which has to quantize its bound + against the sigmoid's qparams. Left alone on int8 codes, ``aten.relu`` + clamps at 0 rather than at the zero point and erases everything below + it -- and it is meant to be a no-op here, sigmoid being positive. + """ + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + return self.conv(torch.relu(torch.sigmoid(x))) + + +class LutActivationClampThenConv2d(torch.nn.Module): + """A two-sided clamp between a LUT activation and a convolution. + + The same gap as ``LutActivationReluThenConv1d``, with bounds of its + own: unquantized, ``aten.clamp`` holds fp32 limits against int8 codes + and keeps an fp32 value on a tensor the convolution reads as int8. + """ + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 1) + + def forward(self, x): + return self.conv(torch.clamp(torch.tanh(x), -0.5, 0.5)) + + +class LutActivationThenElementwise(torch.nn.Module): + """An elementwise binary op fed by a LUT activation. + + ``minimum`` and ``maximum`` pick their kernel from the same fake + tensor the input-side layout permute does, and there is no layout + copy here to involve the validator -- the convolution is what drags + one in. So this is the gap on its own: a graph that is int8 start to + finish and still leaves one op on the portable kernel. + """ + + def __init__(self, fn) -> None: + super().__init__() + self.fn = fn + + def forward(self, x, y): + return self.fn(torch.tanh(x), y) + + class NHWCPaddedConv(torch.nn.Module): def __init__(self, channels, shared_pad): super().__init__() @@ -84,6 +166,12 @@ def _count(exported_program, target) -> int: return sum(node.target == target for node in exported_program.graph.nodes) +def _only(exported_program, target) -> Node: + nodes = [node for node in exported_program.graph.nodes if node.target == target] + assert len(nodes) == 1, f"expected one {target}, found {len(nodes)}" + return nodes[0] + + def _run_explicit_layout_pass_manager(tester: CortexMTester) -> CortexMTester: target_config = CortexMTargetConfig(cpu=CortexM.M55) tester.run_passes( @@ -372,6 +460,156 @@ def test_implementation_transpose_conv2d_strided_pointwise(hardtanh): tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) +def test_explicit_layout_accepts_conv_fed_by_lut_activation(): + """A convolution behind a LUT activation lowers and keeps its numerics.""" + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationThenConv1d().eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.aten.convolution.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +def test_layout_permute_behind_lut_activation_lowers_to_transpose(): + """Every layout move in an int8 chain reaches a CMSIS-NN transpose. + + The convolution has a layout move on each side. The one on its output + is synthesized as ``channels_last.permute_copy``, whose substitution + reads no dtype; the one on its input is an ``aten.permute_copy``, and + its substitution read the same fake tensor the layout validator used + to. Behind a LUT activation that tensor is fp32, so the input-side + move was left on the portable kernel while the output-side move + converted. Numerics are unaffected -- portable ``permute_copy`` + computes the right thing on int8 -- so what this costs is the + invariant: a selective build registering only the kernels the graph + is supposed to need fails to load. + """ + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationThenConv1d().eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize( + "fn,cortex_m_op,aten_op", + [ + (torch.minimum, exir_ops.edge.cortex_m.minimum, exir_ops.edge.aten.minimum), + (torch.maximum, exir_ops.edge.cortex_m.maximum, exir_ops.edge.aten.maximum), + ], +) +def test_elementwise_behind_lut_activation_lowers_to_cortex_m(fn, cortex_m_op, aten_op): + """An elementwise binary op behind a LUT activation reaches CMSIS-NN.""" + torch.manual_seed(0) + inputs = (torch.randn(1, 4, 8), torch.randn(1, 4, 8)) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationThenElementwise(fn).eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, cortex_m_op.default) == 1 + assert _count(program, aten_op.default) == 0 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize( + "model,inputs,bounds", + [ + (LutActivationReluThenConv1d, (torch.randn(1, 8, 16),), (int, type(None))), + (LutActivationClampThenConv2d, (torch.randn(1, 4, 8, 8),), (int, int)), + ], +) +def test_clamp_between_lut_activation_and_conv_is_quantized(model, inputs, bounds): + """A clamp-like activation inside an int8 chain runs on int8 terms. + + The convolution reaches a CMSIS-NN kernel either way. What the fake + tensor costs is the clamp between it and the LUT activation: left in + float, its bounds mean nothing against int8 codes and the tensor the + convolution reads is still declared fp32. + """ + torch.manual_seed(0) + tester = _run_explicit_layout_passes(CortexMTester(model().eval(), inputs)) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + clamp = _only(program, exir_ops.edge.aten.clamp.default) + assert tuple(type(bound) for bound in clamp.args[1:]) == bounds + assert clamp.meta["val"].dtype == torch.int8 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize( + "model,inputs,bounds", + [ + (LutActivationReluThenConv1d, (torch.randn(1, 8, 16),), (int, type(None))), + (LutActivationClampThenConv2d, (torch.randn(1, 4, 8, 8),), (int, int)), + ], +) +def test_clamp_between_lut_activation_and_conv_is_quantized_on_legacy_passes( + model, inputs, bounds +): + """The same clamp on the default pass list, which has no validator. + + Here the clamp is the only thing standing between the LUT activation + and the rest of the graph: nothing rejects an unquantized one, so it + reaches the program. `aten.relu` left alone cuts at code 0 while the + zero point is -128, and `aten.clamp` left on float bounds promotes + the tensor back to fp32 in front of a dequantize that reads int8. + """ + torch.manual_seed(0) + tester = CortexMTester(model().eval(), inputs) + tester.quantize().export().to_edge().run_passes() + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + assert _count(program, exir_ops.edge.aten.relu.default) == 0 + clamp = _only(program, exir_ops.edge.aten.clamp.default) + assert tuple(type(bound) for bound in clamp.args[1:]) == bounds + assert clamp.meta["val"].dtype == torch.int8 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +def test_explicit_layout_accepts_conv_fed_by_view_of_lut_activation(): + """A dtype-agnostic op between the LUT activation and the conv stays int8.""" + torch.manual_seed(0) + inputs = (torch.randn(1, 128),) + tester = _run_explicit_layout_passes( + CortexMTester(LutActivationViewThenConv1d().eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_activation.default) == 1 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +def test_lowered_output_dtype_reads_first_output_of_multi_output_node(): + node = torch.fx.Graph().call_function( + exir_ops.edge.aten.max_pool2d_with_indices.default + ) + node.meta["val"] = ( + torch.empty(1, dtype=torch.int8), + torch.empty(1, dtype=torch.int64), + ) + assert lowered_output_dtype(node) == torch.int8 + + +def test_lowered_output_dtype_declines_a_non_node_argument(): + """Substitutions read it straight off ``node.args``, which holds literals.""" + assert lowered_output_dtype(5.0) is None + + def test_explicit_layout_rejects_unsupported_spatial_operator(): tester = CortexMTester(UnsupportedAvgPool(), (torch.randn(1, 3, 8, 8),)) diff --git a/backends/cortex_m/test/test_propagate_qparams.py b/backends/cortex_m/test/test_propagate_qparams.py new file mode 100644 index 00000000000..209e27ae922 --- /dev/null +++ b/backends/cortex_m/test/test_propagate_qparams.py @@ -0,0 +1,110 @@ +# 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 unittest + +import torch +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.cortex_m.passes.propagate_qparams_pass import ( + PropagateQParamsPass, +) +from executorch.backends.cortex_m.utils import lowered_output_dtype +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Graph, GraphModule, Node + + +_INT8 = QuantArgs(scale=0.5, zp=-3, qmin=-128, qmax=127, dtype=torch.int8) +_INT16 = QuantArgs(scale=0.25, zp=0, qmin=-32768, qmax=32767, dtype=torch.int16) + +_VIEW_COPY = exir_ops.edge.aten.view_copy.default +_PERMUTE_COPY = exir_ops.edge.aten.permute_copy.default +_LAYOUT_PERMUTE_COPY = exir_ops.edge.channels_last.permute_copy.default +_UNSQUEEZE_COPY = exir_ops.edge.aten.unsqueeze_copy.default +_ADD = exir_ops.edge.aten.add.Tensor + + +def _fake(shape, dtype=torch.float32) -> torch.Tensor: + return torch.empty(shape, dtype=dtype) + + +class PropagateQParamsPassTest(unittest.TestCase): + """The pass gives rearranging ops the qparams of the tensor they move. + + Each graph is built by hand rather than exported: the condition under test + is a combination of node target and meta that only arises mid-pipeline, + after FoldAndAnnotateQParamsPass has annotated the computing ops and left + the plumbing with the pre-quantization fake tensor. + """ + + def setUp(self) -> None: + self.graph = Graph() + self.x: Node = self.graph.placeholder("x") + self.x.meta["val"] = _fake((1, 4, 8), torch.int8) + + def _call(self, target, *args, val=None, qparams=None) -> Node: + node = self.graph.call_function(target, args) + node.meta["val"] = _fake((1, 4, 8)) if val is None else val + if qparams is not None: + node.meta["output_qparams"] = {0: qparams} + return node + + def _run(self) -> bool: + self.graph.output(()) + module = GraphModule(torch.nn.Module(), self.graph) + return PropagateQParamsPass().call(module).modified + + def test_a_view_inherits_its_producers_qparams(self) -> None: + producer = self._call(_ADD, self.x, self.x, qparams=_INT8) + view = self._call(_VIEW_COPY, producer, [1, 32]) + + self.assertTrue(self._run()) + self.assertEqual(view.meta["output_qparams"][0], _INT8) + # The fake tensor is the stale fp32 one the trace left; the point of the + # pass is that the dtype no longer has to be read off it. + self.assertEqual(view.meta["val"].dtype, torch.float32) + self.assertEqual(lowered_output_dtype(view), torch.int8) + + def test_a_chain_of_plumbing_resolves_in_one_pass(self) -> None: + """The CWW shape: conv -> layout copy -> view -> permute -> view. + + Topological order is what makes one sweep enough -- each node reads a + producer the same sweep has already annotated. + """ + producer = self._call(_ADD, self.x, self.x, qparams=_INT8) + layout = self._call(_LAYOUT_PERMUTE_COPY, producer, [0, 2, 1]) + view = self._call(_VIEW_COPY, layout, [1, 32]) + permute = self._call(_PERMUTE_COPY, view, [0, 2, 1]) + tail = self._call(_UNSQUEEZE_COPY, permute, 0) + + self.assertTrue(self._run()) + for node in (layout, view, permute, tail): + self.assertEqual( + lowered_output_dtype(node), torch.int8, f"{node.name} not resolved" + ) + + def test_existing_qparams_are_not_overwritten(self) -> None: + producer = self._call(_ADD, self.x, self.x, qparams=_INT8) + view = self._call(_VIEW_COPY, producer, [1, 32], qparams=_INT16) + + self.assertFalse(self._run()) + self.assertEqual(view.meta["output_qparams"][0], _INT16) + + def test_an_unquantized_producer_propagates_nothing(self) -> None: + """A view in an fp32 stretch stays fp32 -- there is nothing to inherit.""" + producer = self._call(_ADD, self.x, self.x) + view = self._call(_VIEW_COPY, producer, [1, 32]) + + self.assertFalse(self._run()) + self.assertNotIn("output_qparams", view.meta) + self.assertEqual(lowered_output_dtype(view), torch.float32) + + def test_a_computing_op_is_not_transparent(self) -> None: + """Only rearranging is dtype-preserving; an add is free to requantize.""" + producer = self._call(_ADD, self.x, self.x, qparams=_INT8) + consumer = self._call(_ADD, producer, producer) + + self.assertFalse(self._run()) + self.assertNotIn("output_qparams", consumer.meta) diff --git a/backends/cortex_m/utils.py b/backends/cortex_m/utils.py index b8d9134c0a4..3dbb5c4dca2 100644 --- a/backends/cortex_m/utils.py +++ b/backends/cortex_m/utils.py @@ -14,6 +14,7 @@ from executorch.exir.dialects.edge._ops import EdgeOpOverload from torch.fx import Node +from torch.fx.node import Argument # L-shift value used in CMSIS-NN for int8 operations SHIFT_INT8 = 20 @@ -198,6 +199,25 @@ def skips_quantized_max_pool2d(node: Node) -> bool: ) +def lowered_output_dtype(node: Argument) -> torch.dtype | None: + """The dtype of this node's first output on device. + + Folded output qparams win over the fake tensor, which stays fp32 on + ops such as tanh and softmax until AtenToCortexMPass rewrites them. + An argument that is not a node, or carries neither, reports None, so + a caller gating a substitution on int8 declines it rather than raising. + """ + if not isinstance(node, Node): + return None + qparams = node.meta.get("output_qparams", {}).get(0) + if qparams is not None: + return qparams.dtype + value = node.meta.get("val") + if isinstance(value, (tuple, list)): + value = value[0] if value else None + return value.dtype if isinstance(value, torch.Tensor) else None + + def is_qualified_int8_node(args) -> bool: try: if len(args) < 6: