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

Summary:

Q/DQ folding (`FoldAndAnnotateQParamsPass`) always folds the quantize/dequantize pair around a constant-producing op (`full`, `full_like`, `arange`, `eye`, `linspace`, `scalar_tensor`) into the op: it removes the Q node and records the output qparams, but the op itself keeps producing float. The Arm backend follows folding with `ComputeConstantOpsAOTPass`, which evaluates such ops in float and quantizes the result. Cortex-M did not, so in `x + torch.full(shape, 0.5)` the int8 `quantized_add` kernel read a float32 constant, and the output was silently wrong.

Add `ComputeQuantizedConstantOpsPass` and run it right after folding on both Cortex-M pass lists. 

- a. It only touches nodes with output qparams, so unquantized constant ops are unchanged. A quantized `full` stays a runtime fill, with its fill value quantized and its dtype set to the qparams' dtype, so a large fill does not turn into a large constant; 
- b.`full_like` is first converted to `full` with the Arm backend's `ConvertFullLikeToFullPass`. 
- c. Since `ComputeConstantOpsAOTPass` only logs a failed precompute, the new pass raises if a quantized constant op is still in the graph afterwards, as it would feed float data to an int8 kernel.

Differential Revision: D128002647

Author: pssrawat

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

@@ -48,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
@@ -149,6 +150,7 @@ class CortexMPassManager(ExportedProgramPassManager):
         # Retrace scalar ops while the graph still has floating-point Q/DQ inputs.
         ReplaceScalarWithTensorArgPass,
         FoldAndAnnotateQParamsPass,
+        ComputeQuantizedConstantOpsPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,
         DecomposeHardswishPass,
@@ -164,6 +166,7 @@ class CortexMPassManager(ExportedProgramPassManager):
     explicit_layout_pass_list: list[PassClass] = [
         RemoveGetItemPass,
         FoldAndAnnotateQParamsPass,
+        ComputeQuantizedConstantOpsPass,
         ReplaceScalarWithTensorArgPass,
         ActivationFusionPass,
         QuantizedClampActivationPass,

backends/cortex_m/test/test_explicit_layout_pipeline.py

@@ -9,6 +9,7 @@
 
 import pytest
 import torch
+from executorch.backends.arm._passes import ComputeConstantOpsAOTPass
 from executorch.backends.cortex_m.passes.cortex_m_pass_manager import (
     CortexMPassManager,
     LiftConstantTensorsPass,
@@ -552,14 +553,100 @@ def test_convs_around_a_reshape_behind_lut_activation_lower_to_cmsis_nn():
     tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
 
 
-def test_explicit_layout_rejects_float_input_to_int8_conv():
+def test_explicit_layout_lowers_conv_on_a_quantized_constant():
+    inputs = (torch.randn(1, 8, 16),)
     model = ThenConv1d(lambda x: torch.full_like(x, 0.5, dtype=torch.float32))
-    tester = CortexMTester(model.eval(), (torch.randn(1, 8, 16),))
+    tester = _run_explicit_layout_passes(CortexMTester(model.eval(), inputs))
+    program = tester.get_artifact(StageType.RUN_PASSES).exported_program()
+
+    assert _count(program, exir_ops.edge.aten.full_like.default) == 0
+    assert _count(program, exir_ops.edge.aten.full.default) == 1
+    assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+@pytest.mark.parametrize(
+    "constant,constant_op,shape",
+    [
+        (
+            lambda: torch.arange(-1.0, 1.0, 0.125),
+            exir_ops.edge.aten.arange.start_step,
+            (16,),
+        ),
+        (
+            lambda: torch.scalar_tensor(0.5),
+            exir_ops.edge.aten.scalar_tensor.default,
+            (),
+        ),
+    ],
+    ids=["arange", "scalar_tensor"],
+)
+def test_quantized_constant_is_computed_before_quantizing(
+    use_explicit_layout, constant, constant_op, shape
+):
+    torch.manual_seed(0)
+    inputs = (torch.randn(shape),)
+    tester = CortexMTester(Apply(lambda x: x + constant()).eval(), inputs)
+    tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout)))
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    assert _count(program, constant_op) == 0
+    assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+@pytest.mark.parametrize(
+    "fill",
+    [lambda x: torch.full((1, 8, 16), 0.5), lambda x: torch.full_like(x, 0.5)],
+    ids=["full", "full_like"],
+)
+def test_quantized_fill_stays_a_runtime_int8_full(use_explicit_layout, fill):
+    torch.manual_seed(0)
+    inputs = (torch.randn(1, 8, 16),)
+    tester = CortexMTester(Apply(lambda x: x + fill(x)).eval(), inputs)
+    tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=use_explicit_layout)))
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    (full,) = program.graph.find_nodes(
+        op="call_function", target=exir_ops.edge.aten.full.default
+    )
+    assert full.meta["val"].dtype == torch.int8
+    assert _count(program, exir_ops.edge.aten.full_like.default) == 0
+    assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+
+
+def test_failed_constant_precompute_is_an_error(monkeypatch):
+    def fail(self, node):
+        raise ValueError("cannot evaluate")
+
+    monkeypatch.setattr(ComputeConstantOpsAOTPass, "compute_node_aot", fail)
+    tester = CortexMTester(
+        Apply(lambda x: x + torch.arange(-1.0, 1.0, 0.125)).eval(),
+        (torch.randn(16),),
+    )
 
     with pytest.raises(Exception) as caught:
         _run_explicit_layout_passes(tester)
 
     error = caught.value
     while error.__cause__ is not None:
         error = error.__cause__
-    assert "must move an int8 tensor" in str(error)
+    assert "Could not precompute quantized constant ops" in str(error)
+
+
+@pytest.mark.parametrize("use_explicit_layout", [False, True])
+def test_unquantized_constant_is_not_precomputed(use_explicit_layout):
+    inputs = (torch.randn(1, 8, 16),)
+    tester = CortexMTester(
+        Apply(lambda x: x + torch.full((1, 8, 16), 0.5)).eval(), inputs
+    )
+    tester.export().to_edge()
+    program = _run_passes(tester, use_explicit_layout)
+
+    assert _count(program, exir_ops.edge.aten.full.default) == 1
+    tester.run_method_and_compare_outputs(inputs=inputs)