XNNPACKQuantizer: apply module filters to ops only, not weight nodes (#24056)

### Summary
Annotators that include weight/bias `get_attr` nodes in the partition
(conv, conv_relu, conv_bn, linear_relu, ...) ran `filter_fn` over every
node. Those nodes have no `nn_module_stack`, so
`set_module_type(torch.nn.Conv2d, ...)` and `set_module_name` never
matched a conv in an exported graph. The filter now checks only the main
op (conv/linear/add/mul), like the linear annotator, so a trailing relu
is fused under its op's config.

### Test plan
`python -m unittest
backends.xnnpack.test.quantizer.test_xnnpack_quantizer
backends.xnnpack.test.quantizer.test_pt2e_quantization` (new
`test_set_module_type_conv2d` and `test_set_module_type_conv2d_relu`
fail without the fix; `test_gru` fails before and after on macOS with
NoQEngine).

cc @GregoryComer @digantdesai @cbilgin @JakeStevens @nil-is-all

Author: msluszniak

backends/xnnpack/quantizer/xnnpack_quantizer_utils.py

@@ -223,7 +223,7 @@ def _annotate_linear_relu(
         if _is_annotated(partition):
             continue
 
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(linear_node):
             continue
 
         linear_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
@@ -310,7 +310,7 @@ def _do_annotate_conv(
         if _is_annotated(partition) or skip:
             continue
 
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(conv_node):
             continue
 
         conv_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
@@ -373,7 +373,7 @@ def _do_annotate_conv_relu(
         if is_conv_transpose and groups != 1:
             continue
 
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(conv_node):
             continue
 
         conv_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
@@ -604,7 +604,7 @@ def _conv_bn(x, conv_weight, conv_bias, bn_weight, bn_bias, bn_rm, bn_rv):
             partition.append(bias_node)
         if _is_annotated(partition):
             continue
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(conv_node):
             continue
 
         # Annotate conv inputs and pattern output
@@ -786,7 +786,7 @@ def _annotate_add_relu(  # noqa: C901
         if _is_annotated(partition):
             continue
 
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(add_node):
             continue
 
         if _is_input_non_float_tensor(add_node):
@@ -845,7 +845,7 @@ def _annotate_add(
         if _is_annotated(partition):
             continue
 
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(add_node):
             continue
 
         if _is_input_non_float_tensor(add_node):
@@ -916,7 +916,7 @@ def _annotate_mul_relu(  # noqa: C901
         if _is_annotated(partition):
             continue
 
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(mul_node):
             continue
 
         if _is_input_non_float_tensor(mul_node):
@@ -975,7 +975,7 @@ def _annotate_mul(
         if _is_annotated(partition):
             continue
 
-        if filter_fn and any(not filter_fn(n) for n in partition):
+        if filter_fn and not filter_fn(mul_node):
             continue
 
         if _is_input_non_float_tensor(mul_node):

backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py

@@ -480,6 +480,75 @@ def forward(self, x):
         ]
         self._test_quantizer(m, example_inputs, quantizer, node_occurrence, node_list)
 
+    def test_set_module_type_conv2d(self):
+        # The weight/bias of a conv are get_attr nodes without nn_module_stack,
+        # so the module type filter must only be checked on the conv itself.
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 8, 3)
+                self.linear = torch.nn.Linear(6, 6)
+
+            def forward(self, x):
+                return self.linear(self.conv(x))
+
+        m = M().eval()
+        example_inputs = (torch.randn(1, 3, 8, 8),)
+        quantizer = XNNPACKQuantizer()
+        quantization_config = get_symmetric_quantization_config(is_per_channel=True)
+        quantizer.set_module_type(torch.nn.Conv2d, quantization_config)
+        node_occurrence = {
+            # conv weight only; the linear is not quantized
+            torch.ops.quantized_decomposed.dequantize_per_channel.default: 1,
+            # input and output of the conv
+            torch.ops.quantized_decomposed.quantize_per_tensor.default: 2,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default: 2,
+        }
+        node_list = [
+            torch.ops.quantized_decomposed.quantize_per_tensor.default,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            torch.ops.aten.conv2d.default,
+            torch.ops.quantized_decomposed.quantize_per_tensor.default,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            torch.ops.aten.linear.default,
+        ]
+        self._test_quantizer(m, example_inputs, quantizer, node_occurrence, node_list)
+
+    def test_set_module_type_conv2d_relu(self):
+        # The relu is called in the parent's forward, so only the conv carries
+        # Conv2d in its nn_module_stack. The relu is fused under the conv's config.
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv2d(3, 8, 3)
+                self.linear = torch.nn.Linear(6, 6)
+
+            def forward(self, x):
+                return self.linear(torch.nn.functional.relu(self.conv(x)))
+
+        m = M().eval()
+        example_inputs = (torch.randn(1, 3, 8, 8),)
+        quantizer = XNNPACKQuantizer()
+        quantization_config = get_symmetric_quantization_config(is_per_channel=True)
+        quantizer.set_module_type(torch.nn.Conv2d, quantization_config)
+        node_occurrence = {
+            # conv weight only; the linear is not quantized
+            torch.ops.quantized_decomposed.dequantize_per_channel.default: 1,
+            # conv input and relu output, nothing between conv and relu
+            torch.ops.quantized_decomposed.quantize_per_tensor.default: 2,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default: 2,
+        }
+        node_list = [
+            torch.ops.quantized_decomposed.quantize_per_tensor.default,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            torch.ops.aten.conv2d.default,
+            torch.ops.aten.relu.default,
+            torch.ops.quantized_decomposed.quantize_per_tensor.default,
+            torch.ops.quantized_decomposed.dequantize_per_tensor.default,
+            torch.ops.aten.linear.default,
+        ]
+        self._test_quantizer(m, example_inputs, quantizer, node_occurrence, node_list)
+
     def test_set_module_name_with_underscores(self) -> None:
         """Test that if a module name has an underscore, we can still quantize it"""