diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index a207d77674d..4b9188909b7 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/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): diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index b66166fedb6..4d3015e32ef 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/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"""