Fix quantized Conv1d same padding with even kernels in XNNPACK (#20734)

SakshamKapoor2911 · JakeStevens · web-flow · commit 10d3009521b0 · 2026-07-16T13:15:17.000-04:00
Fixes #20558. Related to #20553. ## Summary Quantized `nn.Conv1d(..., padding="same")` with an even kernel exports with asymmetric padding (unequal left/right amounts), which cannot be folded into the convolution's symmetric padding field. The original pad-folding approach in this PR was replaced with the explicit-PAD approach from #20553: - **InsertPadQDQPass** — inserts implicit quantize/dequantize pairs after constant_pad_nd nodes in quantized contexts so they serialize as quantized static pads. Refactored with additional guard checks (pad_value, pad_amounts, negative amounts) and correct idempotency. - **ConvolutionConfig._get_act_deps** — pulls zero-valued constant_pad_nd nodes (and their QDQ chain if InsertPadQDQPass already ran) into the convolution's partition for both 1D and 2D convs. The merged method replaces separate 1D and 2D implementations that existed on this branch. - **Regression test** — quantized Conv1d even-kernel same-padding covering both symmetric and asymmetric pad cases, validated by numerical comparison. - Removed unrelated __init__.py re-ordering and no-op conv1d_unsqueeze_pass changes. With these changes, an even-kernel `padding="same"` conv1d graph: ``` dequant -> constant_pad_nd -> convolution ``` becomes (after XNNPACK preprocessing and partitioning): ``` [dequant -> pad -> q -> dq -> conv] (single XNNPACK delegate) ``` ## Test plan ``` python -m pytest backends/xnnpack/test/ops/test_conv1d.py -q # 7 passed (1 new regression test with 4 subTests) python -m pytest backends/xnnpack/test/ops/test_conv2d.py -q # 32 passed python -m pytest backends/xnnpack/test/passes/test_insert_pad_qdq.py -q # 3 passed python -m pytest backends/xnnpack/test/ops/test_static_constant_pad.py -q # 8 passed lintrunner -a # No lint issues ``` cc @GregoryComer @digantdesai @cbilgin @JakeStevens @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --------- Co-authored-by: Jacob Stevens <stevens.jacob1492@gmail.com>
diff --git a/backends/xnnpack/_passes/insert_pad_qdq.py b/backends/xnnpack/_passes/insert_pad_qdq.py
@@ -4,81 +4,85 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from typing import cast, List
+
 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.backends.xnnpack.utils.quant_utils import (
+    is_dequant,
+    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.
+    Inserts implicit quantize/dequantize pairs after constant_pad_nd nodes
+    that sit in a quantized context (input is a dequantize node), so the pad
+    can be serialized as a quantized static pad op.
+
+    Skips pads whose output is already quantized (idempotent).
+
+    Without this pass, a zero-valued constant_pad_nd between a dequantize and
+    a convolution would serialize as fp32 while the conv expects quantized
+    activation, causing a mismatch.
     """
 
-    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):
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
-        for pad in list(graph.nodes):
+        modified = False
+        for node in list(graph.nodes):
             if (
-                pad.op != "call_function"
-                or pad.target != exir_ops.edge.aten.constant_pad_nd.default
+                node.op != "call_function"
+                or node.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
+            pad_input = node.args[0]
+            if not (
+                isinstance(pad_input, torch.fx.Node)
+                and is_dequant(pad_input)
+                and pad_input.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):
+            pad_value = cast(float, node.args[2]) if len(node.args) > 2 else 0.0
+            if pad_value != 0.0:
+                continue
+
+            pad_amounts = cast(List[int], node.args[1])
+            if any(p < 0 for p in pad_amounts):
                 continue
 
-            self._insert_qdq_after(graph, pad, tuple(dq.args[1:]))
+            if any(is_quant(user) for user in node.users):
+                continue
+
+            q_params = pad_input.args[1:]
+
+            with graph.inserting_after(node):
+                q = graph.create_node(
+                    "call_function",
+                    exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
+                    args=(node,) + q_params,
+                )
+                q.meta = node.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)
+
+            node.replace_all_uses_with(dq)
+            q.args = (node,) + q_params
+            modified = True
 
         graph_module.recompile()
-        return PassResult(graph_module, True)
+        return PassResult(graph_module, modified)
diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py
@@ -397,44 +397,48 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
             return False
         return True
 
-    def supported_precision_types(self):
-        return [
-            ConfigPrecisionType.FP32,
-            ConfigPrecisionType.STATIC_QUANT,
-            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)
+        is_transpose = node.args[6]
         if (
             precision != ConfigPrecisionType.FP32
             and not is_transpose
-            and is_2d_conv
+            and is_node(act_input)
             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
+            conv_padding = cast(List[int], node.args[4])
+            is_1d = len(conv_padding) == 1
+            is_2d = len(conv_padding) == 2
+
+            pad_value = (
+                cast(float, act_input.args[2]) if len(act_input.args) > 2 else 0.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:
+            spatial_only = (
+                is_1d
+                and (len(pad_amounts) <= 2 or all(a == 0 for a in pad_amounts[2:]))
+            ) or (
+                is_2d
+                and (len(pad_amounts) <= 4 or all(a == 0 for a in pad_amounts[4:]))
+            )
+
+            if pad_value == 0.0 and all(a >= 0 for a in pad_amounts) 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)
 
+    def supported_precision_types(self):
+        return [
+            ConfigPrecisionType.FP32,
+            ConfigPrecisionType.STATIC_QUANT,
+            ConfigPrecisionType.DYNAMIC_QUANT,
+        ]
+
 
 class AddmmConfig(GEMMConfig):
     """
diff --git a/backends/xnnpack/test/ops/test_conv1d.py b/backends/xnnpack/test/ops/test_conv1d.py
@@ -90,6 +90,24 @@ def forward(self, x):
             z = torch.add(y, z)
             return z
 
+    class Conv1dSamePadding(torch.nn.Module):
+        def __init__(self, kernel_size: int, dilation: int = 1):
+            super().__init__()
+            self.conv1d = torch.nn.Conv1d(
+                in_channels=2,
+                out_channels=4,
+                kernel_size=kernel_size,
+                dilation=dilation,
+                padding="same",
+                bias=True,
+            )
+
+        def forward(self, x):
+            return self.conv1d(x)
+
+    def _get_calibration_samples(self, inputs):
+        return [tuple(torch.randn_like(inputs[i]) for i in range(len(inputs)))]
+
     def _test_conv1d(
         self,
         module,
@@ -102,9 +120,7 @@ def _test_conv1d(
         skip_to_executorch=False,
     ):
         calibration_samples = (
-            [tuple(torch.randn_like(inputs[i]) for i in range(len(inputs)))]
-            if quantized
-            else None
+            self._get_calibration_samples(inputs) if quantized else None
         )
 
         tester = (
@@ -160,6 +176,43 @@ def test_qs8_conv1d(self):
             self.Conv1d(), inputs, 1, quantized=True, dynamic_shape=dynamic_shapes
         )
 
+    def test_qs8_conv1d_even_kernel_same_padding(self):
+        inputs = (torch.randn(1, 2, 16),)
+        configs = [
+            (2, 1),
+            (3, 1),
+            (4, 1),
+            (4, 2),
+        ]
+        for kernel_size, dilation in configs:
+            with self.subTest(kernel_size=kernel_size, dilation=dilation):
+                (
+                    Tester(
+                        self.Conv1dSamePadding(
+                            kernel_size=kernel_size, dilation=dilation
+                        ),
+                        inputs,
+                    )
+                    .quantize(
+                        Quantize(
+                            calibration_samples=self._get_calibration_samples(inputs)
+                        )
+                    )
+                    .export()
+                    .check_count({"torch.ops.aten.conv1d.padding": 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(num_runs=10, atol=0.04, rtol=0.02)
+                )
+
     def test_qs8_conv1d_batchnorm_seq(self):
         inputs = (torch.randn(2, 2, 4),)
         dynamic_shapes = ({0: torch.export.Dim("batch", min=2, max=10)},)
