Only claim conv overloads the QAT fold can actually match

backends/qualcomm/quantizer/conv_bn.py

@@ -47,11 +47,25 @@
 from .qconfig import QuantizationConfig
 from .rules import _is_annotated, _mark_nodes_as_annotated, Q_ANNOTATION_KEY
 
+# Only the overloads torchao's `_fuse_conv_bn_qat` / `_fold_conv_bn_qat` can
+# actually match. Both trace their patterns from `F.conv1d`, `F.conv2d`,
+# `F.conv_transpose1d` and `F.conv_transpose2d`, which export to exactly these
+# four targets.
+#
+# Claiming a conv the fold cannot match is worse than not claiming it: this pass
+# drops the conv's output qspec on the assumption the fold will move the bn
+# output observer onto it, so if the fold misses, the conv is left emitting
+# float32 into a quantized graph and QNN rejects it (`0x232 != 0x408`, FLOAT_32
+# meeting UFIXED_POINT_8). `aten.conv2d.padding` -- produced by
+# `nn.Conv2d(..., padding="same")` -- hit exactly that. `aten.convolution.default`
+# only appears after decomposition, well after this pass runs, so it never
+# matched anything here either.
+#
+# Note `_is_conv_or_conv_transpose_node` is NOT a usable guard: it returns True
+# for `conv2d.padding`.
 CONV_TARGETS = (
     torch.ops.aten.conv1d.default,
     torch.ops.aten.conv2d.default,
-    torch.ops.aten.conv2d.padding,
-    torch.ops.aten.convolution.default,
     torch.ops.aten.conv_transpose1d.default,
     torch.ops.aten.conv_transpose2d.input,
 )

backends/qualcomm/tests/test_conv_bn_qat.py

@@ -65,6 +65,41 @@ def forward(self, x):
         return torch.relu(self.bn2(self.pw(x)))
 
 
+class ConvBnSamePadding(torch.nn.Module):
+    """`padding="same"` exports to `aten.conv2d.padding`, which torchao cannot fold."""
+
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(3, 8, 3, padding="same", bias=False)
+        self.bn = torch.nn.BatchNorm2d(8)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(self.bn(self.conv(x)))
+
+
+class Conv1dBn(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv1d(3, 8, 3, padding=1, bias=False)
+        self.bn = torch.nn.BatchNorm1d(8)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(self.bn(self.conv(x)))
+
+
+class ConvTranspose2dBn(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.ConvTranspose2d(3, 8, 3, padding=1, bias=False)
+        self.bn = torch.nn.BatchNorm2d(8)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(self.bn(self.conv(x)))
+
+
 def _qat_convert(module: torch.nn.Module, example_inputs):
     quantizer = QnnQuantizer()
     quantizer.set_default_quant_config(
@@ -145,6 +180,37 @@ def test_depthwise_separable_conv_bn_folds(self):
         self._assert_conv_biases_quantized(converted)
         self._assert_conv_outputs_quantized(converted)
 
+    def test_conv1d_bn_folds(self):
+        example_inputs = (torch.randn(1, 3, 16),)
+        converted = _qat_convert(Conv1dBn().eval(), example_inputs)
+        self._assert_no_batch_norm(converted)
+        self._assert_conv_outputs_quantized(converted)
+
+    def test_conv_transpose2d_bn_folds(self):
+        example_inputs = (torch.randn(1, 3, 16, 16),)
+        converted = _qat_convert(ConvTranspose2dBn().eval(), example_inputs)
+        self._assert_no_batch_norm(converted)
+        self._assert_conv_outputs_quantized(converted)
+
+    def test_unfoldable_conv_overload_keeps_quantized_output(self):
+        """A conv the fold cannot match must keep its own output qspec.
+
+        `padding="same"` exports to `aten.conv2d.padding`. torchao's
+        `_fuse_conv_bn_qat` / `_fold_conv_bn_qat` trace their patterns from
+        `F.conv1d`/`F.conv2d`/`F.conv_transpose{1,2}d`, so only those four
+        overloads can ever match and BatchNorm survives here.
+
+        That is a torchao limitation this pass cannot fix, but it must not make it
+        worse: if the partition pass claims the conv it drops the conv's output
+        qspec expecting the fold to move the bn output observer onto it, and the
+        conv is left emitting float32 into a quantized graph -- which QNN rejects
+        with `0x232 != 0x408`. So the conv must stay quantized even when the fold
+        misses.
+        """
+        example_inputs = (torch.randn(1, 3, 16, 16),)
+        converted = _qat_convert(ConvBnSamePadding().eval(), example_inputs)
+        self._assert_conv_outputs_quantized(converted)
+
     def test_conv_without_bn_is_still_annotated(self):
         """The partition pass must not starve the per-node conv annotator."""
         example_inputs = (torch.randn(1, 3, 16, 16),)