Fold redundant DQ - > Q (#21016)

Differential Revision: D112564036

Pull Request resolved: #21016

Author: JakeStevens

backends/nxp/edge_passes/fold_redundant_qdq_pass.py

@@ -0,0 +1,36 @@
+# 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.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass
+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.
+
+    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.
+    """
+
+    def run(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        dequant_nodes = [
+            node
+            for node in graph_module.graph.nodes
+            if node.op == "call_function" and node.target in _DEQUANT_OPS
+        ]
+
+        num_nodes_before = len(graph_module.graph.nodes)
+        eliminate_dq_q(graph_module, dequant_nodes)
+        graph_module.graph.eliminate_dead_code()
+        modified = len(graph_module.graph.nodes) != num_nodes_before
+
+        return PassResult(graph_module, modified)

backends/nxp/edge_passes/neutron_edge_pass_manager.py

@@ -6,6 +6,9 @@
 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,
@@ -25,6 +28,7 @@ def __init__(self, passes: list[NeutronEdgePass] = None):
             MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(),
             RemoveUselessAsStridedCopyNodes(),
             ConvertReshapingNodesToViewPass(),
+            FoldRedundantDequantizeQuantizePass(),
         ]
 
         super().__init__(

backends/nxp/tests/BUCK

@@ -106,6 +106,19 @@ fbcode_target(_kind = python_pytest,
     ],
 )
 
+fbcode_target(_kind = python_pytest,
+    name = "test_fold_redundant_qdq",
+    srcs = [
+        "test_fold_redundant_qdq.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/nxp:edge_passes",
+        "//executorch/backends/nxp:neutron_backend",
+        ":executorch_pipeline",
+    ],
+)
+
 fbcode_target(_kind = python_pytest,
     name = "test_convert_1d_conv_to_2d",
     srcs = [

backends/nxp/tests/test_fold_redundant_qdq.py

@@ -0,0 +1,86 @@
+# 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.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
+
+
+class ConvDropoutConvModule(torch.nn.Module):
+    """Two conv clusters separated by an eval-mode dropout (an identity).
+
+    In eval mode the dropout is a no-op, but the quantizer still wraps it in a
+    shared-spec ``dequantize -> quantize`` cluster (identical scale/zero-point on
+    both sides). Lowering eliminates the dropout itself, yet leaves that wrapping
+    quantization behind as a redundant ``dequantize -> quantize`` pair with no
+    compute node between the two conv clusters. The Neutron partitioner cannot
+    delegate a compute-free QDQ island, so it splits ``conv1`` and ``conv2`` into
+    two separate subgraphs -- the same split that
+    ``FoldRedundantDequantizeQuantizePass`` was written to prevent. Folding that
+    pair away lets both convs delegate as a single subgraph.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.conv1 = torch.nn.Conv2d(4, 8, kernel_size=2, bias=False)
+        self.dropout = torch.nn.Dropout(0.5)
+        self.conv2 = torch.nn.Conv2d(8, 8, kernel_size=2, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        x = self.dropout(x)
+        x = self.conv2(x)
+        return x
+
+
+def _count_delegates(edge_program) -> int:
+    graph = edge_program.exported_program().graph_module.graph
+    return sum(
+        1
+        for node in graph.nodes
+        if node.op == "call_function" and node.target == ExecutorchDelegateCall
+    )
+
+
+INPUT_SHAPE = (1, 4, 8, 8)
+
+
+def test_fold_pass_present_merges_into_single_delegate():
+    # The fold pass is part of the default NeutronEdgePassManager.
+    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,)