Precompute quantized constant ops after Q/DQ folding (#24258)

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

backends/cortex_m/passes/compute_quantized_constant_ops_pass.py

@@ -0,0 +1,63 @@
+# 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,
+    ConvertFullLikeToFullPass,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
+from torch.fx import GraphModule
+
+
+def _is_quantized_fill(node: torch.fx.Node) -> bool:
+    return (
+        node.target == exir_ops.edge.aten.full.default
+        and node.kwargs.get("dtype") == node.meta["output_qparams"][0].dtype
+    )
+
+
+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.
+    ``full_like`` is converted to ``full`` first and handled the same way.
+    """
+
+    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 = qparams.quantize_value(torch.tensor(float(node.args[1])))
+            node.args = (node.args[0], int(fill.item()))
+            node.kwargs = {"dtype": qparams.dtype}
+            return True
+        return super().compute_node_aot(node)
+
+    def call(self, graph_module: GraphModule) -> PassResult:
+        converted = ConvertFullLikeToFullPass()(graph_module)
+        result = super().call(converted.graph_module)
+        # The base pass only logs a failure; a quantized constant left behind
+        # would feed float data to an int8 kernel without any error.
+        remaining = [
+            node.name
+            for node in result.graph_module.graph.nodes
+            if node.op == "call_function"
+            and node.target in self.targeted_ops
+            and node.meta.get("output_qparams")
+            and not _is_quantized_fill(node)
+        ]
+        if remaining:
+            raise RuntimeError(
+                f"Could not precompute quantized constant ops: {remaining}"
+            )
+        return PassResult(result.graph_module, converted.modified or result.modified)

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -1,3 +1,3 @@
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -8,12 +8,15 @@
 import inspect
 from typing import Any, Optional, Type
 
+import torch
 from executorch.backends.arm._passes import (
     DeduplicateGetAttrPass,
     FoldAndAnnotateQParamsPass,
     ScalarsToAttributePass,
 )
 from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
+from executorch.backends.cortex_m.utils import _ACTIVATION_FNS
+from executorch.backends.transforms.aten_to_dialect_pass import AtenToDialectPass
 from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import (
     ConvertConv1dToConv2dPass,
 )
@@ -30,10 +33,12 @@
 from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import (
     ReplaceSqueezeAndUnsqueezeWithViewPass,
 )
+from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import (
     ExportedProgramPassBase,
     ExportedProgramPassResult,
     ExportPass,
+    ProxyValue,
 )
 from executorch.exir.pass_manager import ExportedProgramPassManager, PassType
 from executorch.exir.program._program import _transform, lift_constant_tensor_pass
@@ -43,6 +48,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
@@ -82,6 +88,36 @@
         )
 
 
+# Once Q/DQ is folded these ops take int8, but their ATen fake kernels promote
+# it to float, so each retrace would label them and everything behind them float.
+_FLOAT_PROMOTING_OPS = frozenset({*_ACTIVATION_FNS, exir_ops.edge.aten.div.Tensor})
+
+
+class _PreserveQuantizedDtype(ExportPass):
+    def call_operator(self, op, args, kwargs, meta) -> ProxyValue:
+        result = super().call_operator(op, args, kwargs, meta)
+        input_qparams = meta.data.get("input_qparams")
+        output_qparams = meta.data.get("output_qparams")
+        # Match AtenToCortexMPass, which only lowers these ops with both qparams.
+        if op not in _FLOAT_PROMOTING_OPS or not input_qparams or not output_qparams:
+            return result
+        dtype = next(iter(output_qparams.values())).dtype
+        data = torch.empty_like(result.data, dtype=dtype)
+        self.tracer.set_metadata(result.node, data)
+        return ProxyValue(data, result.proxy)
+
+
+def _preserving_quantized_dtype(pass_cls: PassClass) -> PassClass:
+    # Shared passes retrace too, so every pass needs this, except dialect
+    # lowering passes: they lower these ops themselves, and subclassing one
+    # would give it an empty substitution registry.
+    if not issubclass(pass_cls, ExportPass) or issubclass(
+        pass_cls, AtenToDialectPass
+    ):
+        return pass_cls
+    return type(pass_cls.__name__, (_PreserveQuantizedDtype, pass_cls), {})
+
+
 class _CortexMLoweringPass(ExportedProgramPassBase):
     def __init__(
         self, pass_classes: list[PassClass], target_config: CortexMTargetConfig
@@ -99,7 +135,7 @@
             if "target_config" in signature.parameters:
                 kwargs["target_config"] = self.target_config
 
-            transform_pass = pass_cls(**kwargs)
+            transform_pass = _preserving_quantized_dtype(pass_cls)(**kwargs)
             transformed = _transform(exported_program, transform_pass)
             modified |= transformed is not exported_program
             exported_program = transformed
@@ -114,6 +150,7 @@
         # Retrace scalar ops while the graph still has floating-point Q/DQ inputs.
         ReplaceScalarWithTensorArgPass,
         FoldAndAnnotateQParamsPass,
+        ComputeQuantizedConstantOpsPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,
         DecomposeHardswishPass,
@@ -129,6 +166,7 @@
     explicit_layout_pass_list: list[PassClass] = [
         RemoveGetItemPass,
         FoldAndAnnotateQParamsPass,
+        ComputeQuantizedConstantOpsPass,
         ReplaceScalarWithTensorArgPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,