Fold parameter-only subgraphs before XNNPACK partitioning

[john-rocky]

Summary

Fixes #22078.

The XNNPACK partitioner configs require a static weight (is_param_node in partition/config/gemm_configs.py), so a convolution or a linear whose weight is computed from parameters, which is anything under torch.nn.utils.parametrize such as weight_norm and spectral_norm, is declined and left to the portable kernels together with the weight computation. Nothing warns: WhyNoPartition logs at DEBUG, and the model simply runs slow.

This overrides Partitioner.transform_for_pre_decomposition in XnnpackPartitioner to run constant_prop_pass on the ATen program, so the fold is XNNPACK-scoped and nothing changes in to_edge_transform_and_lower's signature. The skip set mirrors the pass's edge-level default: the factory ops that decompose to aten.full, so a scalar fill does not become a stored tensor, and the quantization primitives, so the Q/DQ chain convert_pt2e leaves on a weight stays in place.

The first commit is #22418 (skip impure ops in constant_prop_pass); it is included here so that torch.rand stays in the graph once the pass runs on every XNNPACK export. This PR rebases to one commit once that lands.

Earlier revision of this PR added a constant_prop flag to to_edge_transform_and_lower; reworked after review.

Measurements

wav2vec2-large's positional convolution, weight_norm(Conv1d(1024, 1024, 128, padding=64, groups=16)) at sequence length 499, module alone, synthetic weights, macOS arm64, torch 2.13.0:

	delegates	ops left outside	median latency	max abs diff vs eager
main	2	sum, pow, convolution	3912.1 ms	1.6e-05
this PR	1	none	4.7 ms	1.2e-05

.pte size is the same (33.6 MB) in both cases: the folded weight replaces weight_v, it does not join it.

Test plan

backends/xnnpack/test/test_xnnpack_partitioner.py:

• test_parametrized_weight_is_folded_before_partitioning: a weight_norm Conv1d lowers to a single delegate call with nothing else at the top level, and the runtime output matches eager.
• test_pre_decomposition_folding_keeps_quantization_primitives: on a convert_pt2e graph, the set of quantized_decomposed nodes is identical before and after the hook.

Also run locally with the hook in place: backends/xnnpack/test/ops/test_conv1d.py, backends/xnnpack/test/ops/test_linear.py, exir/program/test/test_program.py, exir/tests/test_passes.py -k constant_prop.

Local run of backends/xnnpack/test/ops/test_linear.py on this machine (main's Python over the 1.4.0 runtime wheel): 58 of 61 pass. The three that fail (qd8_f16/f32_per_channel_int4, qd8_bf16_per_token_weight_per_channel_group_int4) fail identically with the hook removed, so they are the local runtime, not this change.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 15 Awaiting Approval

As of commit b65d16f with merge base 448fbfe ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[JakeStevens]

a few thoughts:

(1) seems there is a bug in the underlying pass for non-determinant ops:

import torch
from executorch.exir import to_edge_transform_and_lower


class RandomAdd(torch.nn.Module):
    def forward(self, x):
      return x + torch.rand(4)


x = torch.zeros(4)
for fold in (False, True):
      ep = torch.export.export(RandomAdd(), (x,), strict=True)
      edge = to_edge_transform_and_lower(ep, constant_prop=fold)
      model = edge.exported_program().module()

      first = model(x)
      second = model(x)

      print(f"{fold=}, outputs_equal={torch.equal(first, second)}")
      print(first)
      print(second)

(2) I don't really like adding a flag to to_edge_transform_and_lower for a pass. maybe better to enable it only for xnnpack in the pre annotation transform. else there may be a better place to hook this in. thoughts maybe @JacobSzwejbka ?