Promote NXP redundant Q/DQ cleanup to to_edge (#22913)

Eval-mode dropout can disappear before shared no-op cleanup runs,
leaving adjacent dequantize/quantize conversions. Promote NXP's
`FoldRedundantDequantizeQuantizePass` from #21016 into the shared
lowering sequence used by `to_edge` and `to_edge_transform_and_lower`.
Cortex-M MLPerf Tiny DS-CNN then finishes with one quantize and one
dequantize operation in both layouts, down from three of each.
MobileNetV2 similarly drops from two to one of each.

The shared `eliminate_dq_q` matcher remains unchanged, keeping this
release patch focused on promoting the existing cleanup. Remove the
stale dropout target from `RemoveNoopPass`. NXP retains a later
invocation of the shared pass because its auxiliary-operator splitting
can introduce new adjacent pairs. Update the Cortex-M model counts and
correct the clone regression's obsolete Q/DQ signatures so the test also
checks numerical equivalence.

Validation on PyTorch 2.14 CPU: 77 EXIR and clone-transform tests
passed, with 28 passing subtests and one expected failure. The dropout
regression covers both edge APIs, per-tensor and per-channel
quantization, and int8/uint8/int16/int32. All 25 selected Cortex-M
DS-CNN, MobileNetV2, and explicit-layout tests passed, including
Corstone-300 FVP implementation tests. Additional explicit-layout DS-CNN
and legacy `ExirExportedProgram.to_edge` numerical checks passed.
Formatting, flake8, and `git diff --check` passed.

Authored with AI assistance from OpenAI Codex.

Co-authored-by: RJ Ascani <rja@meta.com>

Author: pytorchbot

backends/cortex_m/test/models/test_ds_cnn.py

@@ -16,15 +16,15 @@
     "executorch_exir_dialects_edge__ops_aten_relu_default": 9,
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
     "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 18,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 17,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 15,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 15,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 13,
 }
 
 ops_after_transforms: dict[str, int] = {
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
-    "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 3,
+    "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1,
     "executorch_exir_dialects_edge__ops_cortex_m_pad_default": 1,
-    "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 3,
+    "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 4,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 5,

backends/cortex_m/test/models/test_mobilenet_v2.py

@@ -21,14 +21,14 @@
     "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
     "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 104,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 79,
-    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 67,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 78,
+    "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 66,
 }
 
 ops_after_transforms: dict[str, int] = {
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
-    "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 2,
-    "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2,
+    "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1,
+    "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 10,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 35,

backends/nxp/BUCK

@@ -31,6 +31,7 @@ fbcode_target(_kind = runtime.python_library,
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/exir:pass_manager",
+        "//executorch/exir/passes:fold_redundant_qdq_pass",
     ],
 )
 

backends/nxp/edge_passes/neutron_edge_pass_manager.py

@@ -6,28 +6,29 @@
 from executorch.backends.nxp.edge_passes.convert_reshaping_nodes_to_view import (
     ConvertReshapingNodesToViewPass,
 )
-from executorch.backends.nxp.edge_passes.fold_redundant_qdq_pass import (
-    FoldRedundantDequantizeQuantizePass,
-)
 from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import (
     MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass,
     MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass,
 )
-from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass
 from executorch.backends.nxp.edge_passes.remove_as_strided_copy_nodes import (
     RemoveUselessAsStridedCopyNodes,
 )
-from torch.fx.passes.infra.pass_manager import PassManager
+from executorch.exir.pass_base import ExportPass
+from executorch.exir.pass_manager import PassManager
+from executorch.exir.passes.fold_redundant_qdq_pass import (
+    FoldRedundantDequantizeQuantizePass,
+)
 
 
 class NeutronEdgePassManager(PassManager):
 
-    def __init__(self, passes: list[NeutronEdgePass] = None):
-        passes: list[NeutronEdgePass] = passes or [
+    def __init__(self, passes: list[ExportPass] = None):
+        passes: list[ExportPass] = passes or [
             MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(),
             MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(),
             RemoveUselessAsStridedCopyNodes(),
             ConvertReshapingNodesToViewPass(),
+            # Auxiliary-op splitting can introduce DQ -> Q on fanout branches.
             FoldRedundantDequantizeQuantizePass(),
         ]
 

backends/nxp/tests/test_fold_redundant_qdq.py

@@ -6,9 +6,6 @@
 
 import torch
 
-from executorch.backends.nxp.edge_passes.fold_redundant_qdq_pass import (
-    FoldRedundantDequantizeQuantizePass,
-)
 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
 
 ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate
@@ -54,33 +51,10 @@ def _count_delegates(edge_program) -> int:
 
 
 def test_fold_pass_present_merges_into_single_delegate():
-    # The fold pass is part of the default NeutronEdgePassManager.
+    # The fold pass runs during to_edge, before Neutron partitioning.
     edge_program = to_quantized_edge_program(ConvDropoutConvModule(), INPUT_SHAPE)
 
     num_delegates = _count_delegates(edge_program)
     assert (
         num_delegates == 1
     ), f"expected a single delegate with the fold pass, got {num_delegates}"
-
-
-def test_fold_pass_removes_redundant_qdq():
-    graph = torch.fx.Graph()
-    quantized_input = graph.placeholder("quantized_input")
-    qparams = (0.25, 3, -128, 127, torch.int8)
-    dequantize = graph.call_function(
-        torch.ops.quantized_decomposed.dequantize_per_tensor.default,
-        args=(quantized_input, *qparams),
-    )
-    quantize = graph.call_function(
-        torch.ops.quantized_decomposed.quantize_per_tensor.default,
-        args=(dequantize, *qparams),
-    )
-    graph.output(quantize)
-    graph_module = torch.fx.GraphModule(torch.nn.Module(), graph)
-
-    result = FoldRedundantDequantizeQuantizePass().run(graph_module)
-
-    assert result.modified
-    remaining_nodes = list(result.graph_module.graph.nodes)
-    assert [node.op for node in remaining_nodes] == ["placeholder", "output"]
-    assert remaining_nodes[-1].args == (quantized_input,)

backends/transforms/test/test_remove_clone_ops.py

@@ -37,7 +37,7 @@ def test_dq_clone_q_linear(self):
         Test RemoveCloneOpsTransform on a graph with d/q -> clone -> q -> linear pattern
 
         Before: Should contain all nodes
-        After: Should only have the linear operation
+        After: Should only have the final dequantize and linear operations
         """
 
         # Create a graph module directly with the pattern: quant -> clone -> dequant -> fp linear
@@ -60,19 +60,13 @@ def forward(self, x):
         input_node = graph.placeholder("x")
 
         # Create nodes for our pattern: quant -> clone -> dequant -> fp linear
-        # Constants for quantization parameters
-        scale = graph.create_node(
-            "call_function", torch.tensor, args=([0.1],), kwargs={}
-        )
-        zero_point = graph.create_node(
-            "call_function", torch.tensor, args=([0],), kwargs={}
-        )
+        qparams = (0.1, 0, -128, 127, torch.int8)
 
         # Dequantize node
         dequant_node = graph.create_node(
             "call_function",
             torch.ops.quantized_decomposed.dequantize_per_tensor.default,
-            args=(input_node, scale, zero_point, torch.int8),
+            args=(input_node, *qparams),
             kwargs={},
         )
 
@@ -89,19 +83,20 @@ def forward(self, x):
         quant_node = graph.create_node(
             "call_function",
             torch.ops.quantized_decomposed.quantize_per_tensor.default,
-            args=(clone_node, scale, zero_point, torch.int8),
+            args=(clone_node, *qparams),
             kwargs={},
         )
 
-        # Linear node (using the module's linear layer)
-        # Technically, should use quantized weight and bias
-        # but we are just inspecting graph patterns in this test
+        linear_input = graph.call_function(
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            args=(quant_node, *qparams),
+        )
         weight = graph.create_node("get_attr", "linear.weight")
         bias = graph.create_node("get_attr", "linear.bias")
         linear_node = graph.create_node(
             "call_function",
             torch.nn.functional.linear,
-            args=(quant_node, weight, bias),
+            args=(linear_input, weight, bias),
             kwargs={},
         )
 
@@ -125,12 +120,16 @@ def forward(self, x):
         )
 
         # Apply the transform
+        inputs = torch.randint(-128, 128, (2, 10), dtype=torch.int8)
+        expected = gm(inputs)
         transformed_gm = RemoveCloneOpsTransform()(gm).graph_module
 
         # Verify the dq -> clone -> q pattern is removed and linear op is still present using FileCheck
         FileCheck().check_not(
             "executorch_exir_dialects_edge__ops_aten_clone_default"
-        ).check_not("quantized_decomposed.dequantize_per_tensor.default").check_not(
+        ).check_count(
+            "quantized_decomposed.dequantize_per_tensor.default", 1, exactly=True
+        ).check_not(
             "quantized_decomposed.quantize_per_tensor.default"
         ).check_count(
             "torch._C._nn.linear",
@@ -139,6 +138,7 @@ def forward(self, x):
         ).run(
             transformed_gm.code
         )
+        torch.testing.assert_close(transformed_gm(inputs), expected, rtol=0, atol=0)
 
     def test_clone_non_identity_survives(self):
         """Verify clone ops that modify memory_format are preserved by RemoveCloneOpsTransform."""

exir/passes/BUCK

@@ -13,6 +13,7 @@ fbcode_target(_kind = runtime.python_library,
         ":convert_constant_dim_order_pass",
         ":debug_handle_generator_pass",
         ":external_constants_pass",
+        ":fold_redundant_qdq_pass",
         ":init_mutable_pass",
         ":insert_write_back_for_buffers_pass",
         ":legalize_portable_dim_order_pass",
@@ -233,6 +234,18 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "fold_redundant_qdq_pass",
+    srcs = [
+        "fold_redundant_qdq_pass.py",
+    ],
+    deps = [
+        ":remove_noop_pass",
+        "//caffe2:torch",
+        "//executorch/exir:pass_base",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "remove_noop_pass",
     srcs = [

exir/passes/__init__.py

@@ -37,6 +37,9 @@
 from executorch.exir.passes.debug_handle_generator_pass import DebugHandleGeneratorPass
 
 from executorch.exir.passes.executorch_prim_ops_registry import _EXECUTORCH_SYM_OPS
+from executorch.exir.passes.fold_redundant_qdq_pass import (
+    FoldRedundantDequantizeQuantizePass,
+)
 from executorch.exir.passes.insert_write_back_for_buffers_pass import (
     insert_write_back_for_buffers_pass,
 )
@@ -517,6 +520,7 @@ def dead_code_elimination_pass(graph_module: torch.fx.GraphModule) -> PassResult
             RemoveNoopPass(),
             PruneEmptyTensorsPass(),
             RemoveToCopyPass(),
+            FoldRedundantDequantizeQuantizePass(),
         ]
     ).passes
 )

…p/edge_passes/fold_redundant_qdq_pass.py exir/passes/fold_redundant_qdq_pass.pybackends/nxp/edge_passes/fold_redundant_qdq_pass.py renamed to exir/passes/fold_redundant_qdq_pass.py backends/nxp/edge_passes/fold_redundant_qdq_pass.py renamed to exir/passes/fold_redundant_qdq_pass.py

@@ -6,22 +6,19 @@
 
 import torch
 
-from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass
+from executorch.exir.pass_base import ExportPass
 from executorch.exir.passes.remove_noop_pass import _DEQUANT_OPS, eliminate_dq_q
 from torch.fx.passes.infra.pass_base import PassResult
 
 
-class FoldRedundantDequantizeQuantizePass(NeutronEdgePass):
-    """Fold redundant ``dequantize -> quantize`` pairs with identical qparams.
+class FoldRedundantDequantizeQuantizePass(ExportPass):
+    """Fold adjacent ``dequantize -> quantize`` pairs using the shared qparam matcher.
 
-    A dequantize immediately followed by a quantize at identical qparams is the
-    identity on the already-quantized value, so this pass reuses the shared
-    ``eliminate_dq_q`` helper to rewire each such quantize's consumers to the
-    dequantize's quantized input, removing the island and letting the neighboring
-    clusters delegate as a single subgraph.
+    Decomposition can erase a quantized no-op, such as eval-mode dropout, leaving
+    its surrounding dequantize and quantize nodes adjacent.
     """
 
-    def run(self, graph_module: torch.fx.GraphModule) -> PassResult:
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         dequant_nodes = [
             node
             for node in graph_module.graph.nodes
@@ -32,5 +29,8 @@ def run(self, graph_module: torch.fx.GraphModule) -> PassResult:
         eliminate_dq_q(graph_module, dequant_nodes)
         graph_module.graph.eliminate_dead_code()
         modified = len(graph_module.graph.nodes) != num_nodes_before
+        if modified:
+            graph_module.graph.lint()
+            graph_module.recompile()
 
         return PassResult(graph_module, modified)

exir/passes/remove_noop_pass.py

@@ -61,7 +61,6 @@ def call(self, graph_module: GraphModule) -> PassResult:
 
             if node.target not in (
                 torch.ops.aten.to.dtype,
-                torch.ops.aten.dropout.default,
                 torch.ops.aten.slice_copy.Tensor,
             ):
                 continue