Delegate even-kernel 'same'-padding convs via a quantized static pad (#20553)

JakeStevens · web-flow · commit 420a8ab8a526 · 2026-07-08T21:45:06.000Z
Differential Revision: D109871964 Pull Request resolved: #20553
diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py
@@ -27,6 +27,7 @@
 from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
+from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass
 from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
 from executorch.backends.xnnpack._passes.propagate_custom_meta_pass import (
     PropagateCustomMetaPass,
@@ -81,6 +82,7 @@ def __init__(
                 FuseActivationPass,
                 DecomposeConcatenate,
                 RemoveGetItemPass,
+                InsertPadQDQPass,
                 Conv1dUnsqueezePass,
                 PReLUReshapePass,
                 ChannelsLastTaggedReshapePass,
diff --git a/backends/xnnpack/_passes/insert_pad_qdq.py b/backends/xnnpack/_passes/insert_pad_qdq.py
@@ -0,0 +1,84 @@
+# 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.xnnpack._passes.xnnpack_pass import XNNPACKPass
+from executorch.backends.xnnpack.utils.quant_utils import is_quant, tag_as_implicit_q_dq
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
+
+
+class InsertPadQDQPass(XNNPACKPass):
+    """
+    Completes the quantization of a constant_pad_nd that sits inside a quantized
+    region but was left with an fp32 output.
+
+    An even-kernel 'same'-padding conv decomposes (after quantization) into
+    dequant -> constant_pad_nd -> convolution. Because the pad is introduced by
+    to_edge decomposition -- after the quantizer has run -- it is never annotated,
+    so no quantize follows it and its output would serialize as fp32. The
+    downstream conv would then reject its (now fp32) activation.
+
+    A zero-valued pad preserves quantization, so we insert an implicit
+    quantize -> dequantize pair after the pad, reusing the feeding dequant's
+    params. The pad then delegates as a normal quantized XNNStaticConstantPad and
+    the conv sees a proper dequantized activation. The pad node itself is left in
+    place, so all graph shapes stay consistent through later retracing passes.
+    """
+
+    def _insert_qdq_after(self, graph, pad, q_params):
+        with graph.inserting_after(pad):
+            q = graph.create_node(
+                "call_function",
+                exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+                args=(),
+            )
+            q.meta = pad.meta.copy()
+            tag_as_implicit_q_dq(q)
+        with graph.inserting_after(q):
+            dq = graph.create_node(
+                "call_function",
+                exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
+                args=(q,) + q_params,
+            )
+            dq.meta = q.meta.copy()
+            tag_as_implicit_q_dq(dq)
+        pad.replace_all_uses_with(dq)
+        # Set last so replace_all_uses_with above does not rewrite the quantize's
+        # own input.
+        q.args = (pad,) + q_params
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        graph = graph_module.graph
+        for pad in list(graph.nodes):
+            if (
+                pad.op != "call_function"
+                or pad.target != exir_ops.edge.aten.constant_pad_nd.default
+            ):
+                continue
+
+            # Only per-tensor static activations are handled: _insert_qdq_after
+            # builds quantize_per_tensor.default, so the feeding dequant must have
+            # the matching per-tensor signature (a per-channel/per-token/affine
+            # dequant would supply mismatched args).
+            dq = pad.args[0]
+            if (
+                not isinstance(dq, torch.fx.Node)
+                or dq.target
+                != exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
+            ):
+                continue
+
+            # Skip if the pad's output is already quantized. Requiring *no* user to
+            # be a quantize (rather than merely "not all") avoids double-quantizing
+            # pre-existing quant consumers when the pad has mixed users.
+            if not pad.users or any(is_quant(user) for user in pad.users):
+                continue
+
+            self._insert_qdq_after(graph, pad, tuple(dq.args[1:]))
+
+        graph_module.recompile()
+        return PassResult(graph_module, True)
diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py
@@ -38,6 +38,7 @@
     format_target_name,
 )
 from executorch.exir.backend.utils import WhyNoPartition
+from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
 from torch.fx.passes.utils.source_matcher_utils import (
     get_source_partitions,
@@ -403,6 +404,37 @@ def supported_precision_types(self):
             ConfigPrecisionType.DYNAMIC_QUANT,
         ]
 
+    def _get_act_deps(
+        self, node: torch.fx.Node, ep: ExportedProgram, precision: ConfigPrecisionType
+    ) -> Tuple[bool, List[torch.fx.Node]]:
+        # An even-kernel 'same'-padding conv decomposes into
+        # dequant -> constant_pad_nd -> convolution. Pull the zero-valued spatial
+        # pad into this conv's partition so it delegates as a quantized
+        # XNNStaticConstantPad alongside the conv (InsertPadQDQPass completes the
+        # pad's quantization); otherwise the pad is orphaned and the conv is left
+        # with an unquantized activation.
+        # Only supporting 2D, non-transposed convs
+        is_transpose = node.args[6]
+        is_2d_conv = len(cast(List[int], node.args[4])) == 2
+        act_input = get_input_node(node, self.act_idx)
+        if (
+            precision != ConfigPrecisionType.FP32
+            and not is_transpose
+            and is_2d_conv
+            and act_input.target == exir_ops.edge.aten.constant_pad_nd.default
+            and len(act_input.users) == 1
+            and is_dequant(get_input_node(act_input, 0))
+        ):
+            pad_value = act_input.args[2] if len(act_input.args) > 2 else 0
+            pad_amounts = cast(List[int], act_input.args[1])
+            spatial_only = len(pad_amounts) <= 4 or all(a == 0 for a in pad_amounts[4:])
+            if pad_value == 0 and spatial_only:
+                valid, deps = super()._get_act_deps(act_input, ep, precision)
+                if valid:
+                    return (True, [act_input, *deps])
+
+        return super()._get_act_deps(node, ep, precision)
+
 
 class AddmmConfig(GEMMConfig):
     """
diff --git a/backends/xnnpack/test/ops/test_conv2d.py b/backends/xnnpack/test/ops/test_conv2d.py
@@ -310,6 +310,126 @@ def test_qs8_conv2d_per_channel(self) -> None:
                 quant_config=get_symmetric_quantization_config(is_per_channel=True),
             )
 
+    def test_qs8_conv2d_even_kernel_same_padding(self) -> None:
+        # An even-kernel 'same'-padding conv decomposes into
+        # dequant -> constant_pad_nd -> convolution. The pad and conv are pulled
+        # into one partition and both delegate (the pad as a quantized
+        # XNNStaticConstantPad), so neither survives in the top-level graph.
+        for has_bias in (True, False):
+            m = Conv2d(
+                in_channels=2,
+                out_channels=4,
+                kernel_size=(4, 4),
+                stride=(1, 1),
+                padding="same",
+                bias=has_bias,
+            )
+            tester = Tester(m.eval(), m.get_inputs())
+            tester.quantize(
+                Quantize(quantization_config=get_symmetric_quantization_config())
+            )
+            (
+                tester.export()
+                .check_count({"torch.ops.aten.conv2d": 1})
+                .to_edge_transform_and_lower()
+                .check_not(
+                    [
+                        "executorch_exir_dialects_edge__ops_aten_convolution_default",
+                        "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default",
+                    ]
+                )
+                .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+                .to_executorch()
+                .serialize()
+                .run_method_and_compare_outputs(qtol=2)
+            )
+
+    def test_qs8_conv2d_depthwise_even_kernel_same_padding(self) -> None:
+        # Depthwise is a standard (non-transposed) 2D conv, so its even-'same' pad
+        # is delegated alongside the conv too. Assert it still fully delegates and is
+        # numerically correct.
+        m = Conv2d(
+            in_channels=4,
+            out_channels=4,
+            groups=4,
+            kernel_size=(4, 4),
+            stride=(1, 1),
+            padding="same",
+        )
+        tester = Tester(m.eval(), m.get_inputs())
+        tester.quantize(
+            Quantize(quantization_config=get_symmetric_quantization_config())
+        )
+        (
+            tester.export()
+            .check_count({"torch.ops.aten.conv2d": 1})
+            .to_edge_transform_and_lower()
+            .check_not(
+                [
+                    "executorch_exir_dialects_edge__ops_aten_convolution_default",
+                    "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default",
+                ]
+            )
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(qtol=2)
+        )
+
+    class EvenSamePadFlatten(torch.nn.Module):
+        # Reproduces the failure where an even-kernel 'same' conv is followed by a
+        # shape-dependent op. The flatten bakes a fixed view over the conv's padded
+        # output extent; if the pad's spatial contribution is dropped from the graph
+        # the conv output shrinks and the view becomes invalid.
+        def __init__(self):
+            super().__init__()
+            self.conv = torch.nn.Conv2d(1, 16, (2, 1), padding="same")
+
+        def forward(self, x):
+            out = torch.reshape(x, (1, 1, 256, 1))
+            out = self.conv(out)
+            return torch.flatten(out, 1)
+
+        def get_inputs(self):
+            return (torch.randn(1, 1, 256, 1),)
+
+    def test_fp32_even_kernel_same_padding_flatten(self) -> None:
+        m = self.EvenSamePadFlatten().eval()
+        (
+            Tester(m, m.get_inputs())
+            .export()
+            .to_edge_transform_and_lower()
+            .check_not(
+                [
+                    "executorch_exir_dialects_edge__ops_aten_convolution_default",
+                    "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default",
+                ]
+            )
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
+    def test_qs8_even_kernel_same_padding_flatten(self) -> None:
+        m = self.EvenSamePadFlatten().eval()
+        (
+            Tester(m, m.get_inputs())
+            .quantize(Quantize(quantization_config=get_symmetric_quantization_config()))
+            .export()
+            .to_edge_transform_and_lower()
+            .check_not(
+                [
+                    "executorch_exir_dialects_edge__ops_aten_convolution_default",
+                    "executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default",
+                ]
+            )
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(qtol=2)
+        )
+
     def test_fp32_conv2d_seq(self) -> None:
         for transpose in (True, False):
             self._test(Conv2dSeq(transpose=transpose), conv_count=2)
diff --git a/backends/xnnpack/test/passes/test_insert_pad_qdq.py b/backends/xnnpack/test/passes/test_insert_pad_qdq.py
@@ -0,0 +1,73 @@
+# 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 unittest
+
+import torch
+from executorch.backends.test.harness.stages import StageType
+from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass
+from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
+    get_symmetric_quantization_config,
+)
+from executorch.backends.xnnpack.test.tester import Quantize, RunPasses, Tester
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+class TestInsertPadQDQ(unittest.TestCase):
+    Q = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
+    PAD = exir_ops.edge.aten.constant_pad_nd.default
+
+    class PadConv(torch.nn.Module):
+        # Even kernel + 'same' padding decomposes into constant_pad_nd -> conv.
+        def __init__(self):
+            super().__init__()
+            self.conv = torch.nn.Conv2d(2, 2, (2, 2), padding="same")
+
+        def forward(self, x):
+            return self.conv(x)
+
+    def _run(self, passes, quantize):
+        tester = Tester(self.PadConv().eval(), (torch.randn(1, 2, 8, 8),))
+        if quantize:
+            tester.quantize(
+                Quantize(quantization_config=get_symmetric_quantization_config())
+            )
+        artifact = (
+            tester.export()
+            .to_edge()
+            .run_passes(RunPasses(passes))
+            .get_artifact(StageType.RUN_PASSES)
+        )
+        return artifact.exported_program().graph_module.graph
+
+    def _pad(self, graph):
+        return next(n for n in graph.nodes if n.target == self.PAD)
+
+    def _num_quant(self, graph):
+        return sum(1 for n in graph.nodes if n.target == self.Q)
+
+    def test_inserts_qdq_after_quantized_pad(self):
+        # dequant -> pad -> conv: the pass quantizes the pad's output, so the pad
+        # is now consumed by an inserted quantize.
+        graph = self._run([InsertPadQDQPass], quantize=True)
+        pad = self._pad(graph)
+        self.assertTrue(all(user.target == self.Q for user in pad.users))
+
+    def test_idempotent_skips_already_quantized_pad(self):
+        # A second run must see the pad's user is already a quantize and skip it,
+        # so no extra quantize is inserted.
+        once = self._num_quant(self._run([InsertPadQDQPass], quantize=True))
+        twice = self._num_quant(
+            self._run([InsertPadQDQPass, InsertPadQDQPass], quantize=True)
+        )
+        self.assertEqual(once, twice)
+
+    def test_noop_for_fp32_pad(self):
+        # The pad is not fed by a dequant, so the pass leaves the graph untouched.
+        graph = self._run([InsertPadQDQPass], quantize=False)
+        self.assertEqual(self._num_quant(graph), 0)
+        pad = self._pad(graph)
+        self.assertFalse(any(user.target == self.Q for user in pad.users))
