Fix quantized Conv1d same padding with even kernels in XNNPACK

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)

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):
     """

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)},)