Give the shape and layout plumbing the qparams of the tensor it moves

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",
     ],

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

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

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/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"],