Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nil-is-all merged 2 commits into from
Oct 29, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
49d67c2
XNNPACKQuantizer: apply module filters to ops only, not weight nodes
msluszniak Oct 21, 2026
05d9be3
XNNPACKQuantizer: check module filters on the main op only
msluszniak Oct 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions backends/xnnpack/quantizer/xnnpack_quantizer_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
69 changes: 69 additions & 0 deletions backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"""

Expand Down
Loading