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/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/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/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/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: