diff --git a/backends/qualcomm/quantizer/conv_bn.py b/backends/qualcomm/quantizer/conv_bn.py index 4fc7875ce6c..f54ac577131 100644 --- a/backends/qualcomm/quantizer/conv_bn.py +++ b/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, ) diff --git a/backends/qualcomm/tests/test_conv_bn_qat.py b/backends/qualcomm/tests/test_conv_bn_qat.py index 9f604e55249..1b17e789c60 100644 --- a/backends/qualcomm/tests/test_conv_bn_qat.py +++ b/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),)