Skip to content

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

Merged
nil-is-all merged 2 commits into
pytorch:mainfrom
msluszniak:ms/xnnpack-quantizer-module-filter-conv
Oct 29, 2026
Merged

nil-is-all merged 2 commits into
pytorch:mainfrom
msluszniak:ms/xnnpack-quantizer-module-filter-conv

Conversation

@msluszniak

@msluszniak msluszniak commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

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

@pytorch-bot

pytorch-bot Bot commented Oct 21, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24056

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 05d9be3 with merge base dccea61 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 21, 2026
@executorch-triage executorch-triage Bot added the community: contribution PRs coming from community (excluding hardware partners) label Oct 21, 2026
@nil-is-all nil-is-all added module: xnnpack Issues related to xnnpack delegation and the code under backends/xnnpack/ release notes: xnnpack Changes to the XNNPack backend delegate labels Oct 21, 2026
@JakeStevens

Copy link
Copy Markdown
Contributor

Thanks! Overall this makes sense. One gap remains: conv+relu (and linear+relu) under set_module_type(nn.Conv2d) still never gets quantized. The trailing relu is called in the parent's forward, so its nn_module_stack doesn't include Conv2d, and the filter rejects the conv+relu group.

Could we filter on the main op node only, as _annotate_linear already does? A trailing relu then gets quantized and fused under its conv's config, which is what someone asking to quantize Conv2d expects. That also makes _filtered_out unnecessary.

This would mean that you even if you exclude relu, it can get pulled into a preceding quantized conv. But, I dont think that is a very common case

@msluszniak

Copy link
Copy Markdown
Contributor Author

Done, filter now checks only the main op; added a conv+relu test.

@JakeStevens JakeStevens left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

try rebasing so we can check CI and should be good to go

@msluszniak

Copy link
Copy Markdown
Contributor Author

These fails were infra specific, but rebasing either way.

Annotators that put the weight/bias get_attr nodes into the partition
(conv, conv_relu, conv_bn, linear_relu, ...) ran filter_fn over the whole
partition. Those nodes carry no nn_module_stack, so set_module_type and
set_module_name never matched a conv in an exported graph. Skip get_attr
and placeholder nodes, as the linear annotator already does by checking
only the op.
A relu called in the parent's forward has no Conv2d/Linear in its
nn_module_stack, so set_module_type(nn.Conv2d) rejected every conv+relu
partition. Check filter_fn on the conv/linear/add/mul node alone, as
_annotate_linear does; a trailing relu is then fused under its op's config.
@msluszniak
msluszniak force-pushed the ms/xnnpack-quantizer-module-filter-conv branch from 0ef33af to 05d9be3 Compare October 28, 2026 14:16
@nil-is-all
nil-is-all merged commit 10b5018 into pytorch:main Oct 29, 2026
337 of 338 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. community: contribution PRs coming from community (excluding hardware partners) module: xnnpack Issues related to xnnpack delegation and the code under backends/xnnpack/ release notes: xnnpack Changes to the XNNPack backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants