Substitute on qparams, not the input fake tensor (#24121)

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)
 

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 (

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:

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),))