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.

Two things the hook has to get right at this stage of the pipeline:

• The program is not functionalized yet. A KV-cache update is still an in-place copy_, index_put_ or custom op on the buffer, and the graph signature lists no mutated buffers. constant_prop_pass reads mutation from the signature, so on its own it takes such a buffer for a constant and folds the view that is written to; run_decompositions then fails on the aliasing (expected compiled_fn to be GraphModule), which is what the Voxtral realtime job hit on the previous revision. The hook functionalizes first with run_decompositions({}), the call to_edge_transform_and_lower makes right after it anyway. Cost on a 12-layer, 134M-parameter decoder with 24 cache writes: 0.6 s.
• The skip set. The factory ops that decompose to aten.full (full, new_full, ones, new_ones, zeros, new_zeros) and to aten.full_like (full_like, ones_like, zeros_like), so a scalar fill does not become a stored tensor, for the same reason the pass skips full at the edge level. And every quantized_decomposed / torchao op in the graph, so the Q/DQ chain convert_pt2e or quantize_ leaves on a weight stays in place.

Earlier revisions: a constant_prop flag on to_edge_transform_and_lower (reworked after review), then the hook without the functionalize step (broke the Voxtral export). The impure-op fix the first review found in the pass itself landed as #22418.

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.

Quantized paths checked with and without the hook, Q/DQ op set and .pte size identical on all seven: legacy 8da4w, quantize_ 8da4w / 8da8w / int4 weight-only, PT2E static / dynamic / qc4w.

Voxtral realtime, CI export path (streaming, 8da4w linears, 8w embedding, XNNPACK) on a 48x scaled-down random-weight instance of the same graph: the previous revision fails with the CI assertion, this revision lowers, with the same delegate count per method as main (21 / 13 / 0).

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 and on a groupwise int4 weight, the set of quantized_decomposed nodes is identical before and after the hook.
• test_pre_decomposition_folding_skips_factory_ops: every op in the skip set is checked to decompose to full / full_like and to survive the hook.
• test_pre_decomposition_folding_keeps_mutated_buffer: a buffer read and written in place stays a mutated buffer through the hook and through lowering, and the runtime carries its state across calls like eager.

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.

[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 28837bf with merge base 33ed3c5 ():

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 ?

[JakeStevens]

some feedback:

1. The new test imports XNNPACK quantizer and TorchAO, but test_xnnpack_partitioner does not depend on :xnnpack_quantizer or //pytorch/ao:torchao. Add both to its BUCK target.

2. Factory skip list is incomplete. aten.new_full, aten.new_ones, and aten.new_zeros also decompose to aten.full. Add these overloads and a regression test

after these fixes please make sure to run linter as that is failing already

[john-rocky]

Thanks for the review, @JakeStevens. All three points are in, plus one more fix the CI run turned up. The branch is rebased on main now that #22418 landed: the hook (7dec35b) and the review fixes (28837bf).

Your three points. test_xnnpack_partitioner now depends on :xnnpack_quantizer and //pytorch/ao:torchao. new_full, new_ones and new_zeros are in the skip set, and test_pre_decomposition_folding_skips_factory_ops walks the whole set: each op is checked to decompose to full or full_like and to survive the hook. One thing that check showed: full_like / ones_like / zeros_like decompose to aten.full_like, which the edge-level default does not skip. I kept them, since the size argument is the same; happy to trim the set to the edge default if you prefer. flake8 and ufmt are clean on the changed files (the B020 was mine).

The Voxtral job was this PR, not the trunk. The ATen program at this hook is not functionalized yet: the KV-cache update is still an in-place op on the buffer (llama.update_cache_with_indices in Voxtral, index_put_ / copy_ elsewhere), and the graph signature lists no mutated buffers. constant_prop_pass reads mutation from the signature, so it took k_cache / v_cache for constants and folded views of them; run_decompositions then failed on the aliasing between a folded view and the buffer the custom op writes, with the expected compiled_fn to be GraphModule assertion. The hook now functionalizes first with run_decompositions({}), the call _gen_edge_manager_for_partitioners makes right after it anyway; on a 12-layer, 134M-parameter decoder with 24 cache writes that costs 0.6 s. test_pre_decomposition_folding_keeps_mutated_buffer covers it through the runtime, with the state carried across calls. I did not pull the gated checkpoint, so I ran the CI export path (streaming, 8da4w, 8w embedding, XNNPACK) on a 48x scaled-down random-weight instance of the same graph: the previous revision fails with the CI assertion, this one lowers with the same delegate count per method as main.

On the hook location, from your first note: transform_for_annotation only runs under prepare_pt2e / prepare_qat_pt2e, and the case in #22078 is fp32, so transform_for_pre_decomposition is the one spot every XNNPACK export passes through. If you or @JacobSzwejbka would rather have an off switch on the partitioner, that is a small addition.

Of the other red jobs on the previous run: the three unittest jobs died in sccache (AWS credentials) while building the tokenizers wheel, lintrunner-mypy was the backends/arm errors #22482 has since fixed on main, and moshi / qnn are the flaky and trunk failures Dr. CI flagged.

Thanks again for the careful look.